import uuid
from typing import List, Optional, Any
from bson import ObjectId
from qdrant_client.http.models import PointStruct

from ..repositories import repo
from ..llm_utils import LLM_UTILS
from ..config import Config
from .schema import QAList, FAQIndexRequest
from .prompt import get_FAQ_prompt, get_FAQ_user_prompt
from ..logger import log

# Mongo collection
document_col = repo.collection("documents")
agent_col = repo.collection("agents")

# System-level prompt for FAQ generation
system_prompt = get_FAQ_prompt()


def get_knowledge_base_ids_by_agent(data: FAQIndexRequest) -> List[Any]:
    """Return knowledge_base_id array (as strings) for a given agent_id."""
    log.debug(f"Fetching knowledge_base_ids for agent_id={data.agent_id}")
    agent = agent_col.find_one({"_id": ObjectId(data.agent_id)}, {"knowledge_base_id": 1})
    if not agent or "knowledge_base_id" not in agent:
        log.info(f"No knowledge_base_ids found for agent_id={data.agent_id}")
        return []
    kb_ids = [ObjectId(kb_id) for kb_id in agent["knowledge_base_id"]]
    log.info(f"Found {len(kb_ids)} knowledge_base_ids for agent_id={data.agent_id}")
    return kb_ids


def get_document_content_chunks(
    knowledge_base_ids: List[str], chunk_size: int = 8000
) -> List[str]:
    """
    Fetch 'content' fields of documents by knowledge_base_id(s) and return them in chunks.
    """
    try:
        log.debug(f"Fetching documents for knowledge_base_ids={knowledge_base_ids}")
        cursor = document_col.find(
            {"knowledge_base_id": {"$in": knowledge_base_ids}}, {"content": 1, "_id": 0}
        )

        chunks = []
        for doc in cursor:
            content = doc.get("content", "")
            for i in range(0, len(content), chunk_size):
                chunks.append(content[i : i + chunk_size])

        log.info(
            f"Created {len(chunks)} content chunks for knowledge_base_ids={knowledge_base_ids}"
        )
        return chunks

    except Exception as e:
        log.error(f"Error fetching document chunks: {e}", exc_info=True)
        return []


def get_QA_pairs_llm(content_chunks: List[str]) -> Optional[QAList]:
    """Use LLM to generate structured QA pairs from document content."""
    if not content_chunks:
        log.info("No content chunks provided for QA generation")
        return None

    structured_model = LLM_UTILS.llm.with_structured_output(QAList)
    user_prompt = get_FAQ_user_prompt(content_chunks)
    log.debug("Invoking LLM for QA generation...")
    return structured_model.invoke([system_prompt, user_prompt])


def create_points(agent_id: str, qa_results: QAList) -> List[PointStruct]:
    """Convert QA pairs into Qdrant points with embeddings."""
    if not qa_results or not qa_results.qa_pairs:
        log.info(f"No QA pairs available to create points for agent_id={agent_id}")
        return []

    points = []
    for qa in qa_results.qa_pairs:
        vector = LLM_UTILS.embeddings.embed_query(qa.question, output_dimensionality=768)
        payload = {
            "document_id": agent_id,
            "question": qa.question,
            "answer": qa.answer,
        }
        points.append(PointStruct(id=str(uuid.uuid4()), vector=vector, payload=payload))

    log.info(f"Created {len(points)} Qdrant points for agent_id={agent_id}")
    return points


def index_QA_doc(data: FAQIndexRequest) -> bool:
    """Main entry: fetch doc → generate QA pairs → embed → index in Qdrant."""
    try:
        log.info(f"Starting QA indexing for agent_id={data.agent_id}")
        knowledge_base_ids = get_knowledge_base_ids_by_agent(data.agent_id)
        if not knowledge_base_ids:
            log.info(
                f"No knowledge_base_ids for agent_id={data.agent_id}, skipping indexing"
            )
            return False

        content_chunks = get_document_content_chunks(knowledge_base_ids)
        qa_results = get_QA_pairs_llm(content_chunks)
        if not qa_results:
            log.info(f"No QA generated for agent_id={data.agent_id}")
            return False

        points = create_points(data.agent_id, qa_results)
        if not points:
            log.info(f"No points created for agent_id={data.agent_id}")
            return False

        LLM_UTILS.qdrant.upsert(
            collection_name=Config.RAG_QDRANT_COLLECTION,
            points=points,
        )
        log.info(f"Successfully indexed {len(points)} QA pairs for agent_id={data.agent_id}")
        return True
    except Exception as e:
        log.error(f"Error indexing QA doc for agent_id={data.agent_id}: {e}", exc_info=True)
        return False