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 os
import json

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."""
        collections = self.client.get_collections().collections

        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 = []

        for i, chunk in enumerate(chunks):
            point = PointStruct(
                id=i,
                vector=chunk["embedding"],
                payload={
                    "content": chunk["content"],
                    "metadata": chunk["metadata"],
                }
            )
            points.append(point)

        if points:
            self.client.upsert(
                collection_name=self.collection_name,
                points=points,
            )
            print(f"Added {len(points)} chunks to Qdrant")

    def search(self, query_embedding: List[float], top_k: int = 3) -> List[Dict]:
        """Search for similar chunks."""
        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."""
        points = self.client.scroll(
            collection_name=self.collection_name,
            limit=10000,
        )[0]

        books = set()
        for point in points:
            metadata = point.payload["metadata"]
            books.add(metadata.get("title", "Unknown"))

        return sorted(list(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

        # Setup GPT4All
        from gpt4all import GPT4All

        self.llm = GPT4All(model_name=model_name, device="cpu")

    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ê é um assistente especializado em espiritismo. Use as informações abaixo para responder à pergunta do usuário.

Informações:
{context_text}

Instruções:
- Responda de forma clara e concisa
- Use o contexto fornecido
- Quando citar passagens, mencione o livro e capítulo (se disponível)
- Se não souber a resposta, diga que não encontrou na base de dados"""

        response = self.llm.generate(
            system_prompt,
            query,
            max_tokens=500,
        )

        return response

    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": [chunk["metadata"] for chunk in chunks],
        }