from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct, Filter, FieldCondition, MatchValue
import numpy as np
from typing import List, Dict, Optional
import json
import os
import time
import uuid

class QdrantService:
    def __init__(self, url: str, collection_name: str = "espiritismo_books"):
        self.client = QdrantClient(url=url, prefer_grpc=False)
        self.collection_name = collection_name
        self.embedding_dim = 384

        self._ensure_collection()

    def _ensure_collection(self):
        """Create collection if it doesn't exist."""
        last_error = None
        for attempt in range(15):
            try:
                collections = self.client.get_collections().collections
                break
            except Exception as exc:
                last_error = exc
                print(f"Qdrant indisponível ({attempt + 1}/15): {exc}")
                time.sleep(2)
        else:
            raise RuntimeError(f"Não foi possível conectar ao Qdrant: {last_error}")

        collection_exists = any(
            c.name == self.collection_name for c in collections
        )

        if not collection_exists:
            self.client.create_collection(
                collection_name=self.collection_name,
                vectors_config=VectorParams(
                    size=self.embedding_dim,
                    distance=Distance.COSINE,
                ),
            )
            print(f"Created collection: {self.collection_name}")
        else:
            print(f"Collection {self.collection_name} already exists")

    def add_chunks(self, chunks: List[Dict]):
        """Add chunks to the collection."""
        points = [
            PointStruct(
                id=str(chunk.get("id") or uuid.uuid4()),
                vector=chunk["embedding"],
                payload={
                    "content": chunk["content"],
                    "metadata": chunk["metadata"],
                },
            )
            for chunk in chunks
        ]

        for start in range(0, len(points), 128):
            batch = points[start:start + 128]
            self.client.upsert(collection_name=self.collection_name, points=batch)
        if points:
            print(f"Added {len(points)} chunks to Qdrant")

    def search(self, query_embedding, top_k: int = 3) -> List[Dict]:
        """Search for similar chunks."""
        if hasattr(query_embedding, "tolist"):
            query_embedding = query_embedding.tolist()
        results = self.client.search(
            collection_name=self.collection_name,
            query_vector=query_embedding,
            limit=top_k,
        )

        return [
            {
                "content": hit.payload["content"],
                "metadata": hit.payload["metadata"],
                "score": hit.score,
            }
            for hit in results
        ]

    def get_books(self) -> List[str]:
        """Get list of books in the collection."""
        books = set()
        offset = None
        while True:
            points, offset = self.client.scroll(
                collection_name=self.collection_name,
                limit=256,
                offset=offset,
                with_vectors=False,
            )
            for point in points:
                metadata = (point.payload or {}).get("metadata") or {}
                books.add(metadata.get("title", "Unknown"))
            if offset is None:
                break

        return sorted(books)

    def get_collection_info(self) -> Dict:
        """Get information about the collection."""
        info = self.client.get_collection(self.collection_name)
        return {
            "status": info.status,
            "vectors_count": info.vectors_count,
            "points_count": info.points_count,
        }


class RAGService:
    def __init__(self, qdrant_service: QdrantService, model_name: str = "llama-3.2-3b-instruct-q4_0.gguf"):
        self.qdrant = qdrant_service
        self.model_name = model_name
        self.llm = None

    def _get_llm(self):
        if self.llm is None:
            if not os.path.isfile(self.model_name):
                raise FileNotFoundError(
                    f"Modelo não encontrado em {self.model_name}. "
                    "Baixe llama-3.2-3b-instruct-q4_0.gguf para backend/models/."
                )
            from gpt4all import GPT4All

            model_dir = os.path.dirname(self.model_name) or "."
            model_file = os.path.basename(self.model_name)
            self.llm = GPT4All(
                model_name=model_file,
                model_path=model_dir,
                device="cpu",
                allow_download=False,
            )
        return self.llm

    def generate_response(self, query: str, context: List[str], top_k: int = 3) -> str:
        """Generate a response based on the query and context."""
        context_text = "\n\n".join(context)

        system_prompt = f"""Você responde perguntas sobre espiritismo usando apenas o contexto.
Escreva uma ou duas frases completas, em português. Não repita a pergunta nem as instruções.
Se o contexto não bastar, responda apenas: Não encontrei essa informação nos livros indexados.

Contexto:
{context_text}"""
        template = (
            "<|start_header_id|>user<|end_header_id|>\n\n"
            "{0}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
        )
        with self._get_llm().chat_session(system_prompt=system_prompt, prompt_template=template):
            response = self.llm.generate(query, max_tokens=160, temp=0.1)
        return response.strip()

    def query(self, query: str, top_k: int = 3) -> Dict:
        """Query the RAG system."""
        # Get embeddings
        from app.utils.embeddings import EmbeddingService
        embedding_service = EmbeddingService()

        query_embedding = embedding_service.embed_query(query)

        # Search
        chunks = self.qdrant.search(query_embedding, top_k)

        # Generate response
        context = [chunk["content"] for chunk in chunks]

        response = self.generate_response(query, context, top_k)

        return {
            "query": query,
            "response": response,
            "context": context,
            "sources": [
                {"metadata": chunk["metadata"], "score": chunk["score"]}
                for chunk in chunks
            ],
        }