import tiktoken
from typing import Optional
from fastapi import Request

from qdrant_client import models
from ..vectorstore import VectorStoreFactory
from ..embedding import EmbeddingFactory
from ..llm import LLMFactory
from ...core.logger import SingletonLogger
from .pdf_parser import load_pdf_content
from ...config.prompts import PAPER_SUMMARY_SYSTEM_PROMPT

logger = SingletonLogger().get_logger()


class SummaryEngine:
    """SummaryEngine class for generating summaries of papers."""

    DEFAULT_MODEL = "groq/qwen3-32b"

    @staticmethod
    async def generate_paper_summary(
        arxiv_id: str = None, pdf_url: str = None, request: Optional[Request] = None,
        model_name: Optional[str] = None,
    ) -> str:
        try:
            embedding = EmbeddingFactory.build_embedding_model(request=request)
            vector_store = VectorStoreFactory.build_vector_store(
                embedding_model=embedding
            )
            full_content = ""

            if arxiv_id:
                filter_condition = models.Filter(
                    must=[
                        models.FieldCondition(
                            key="metadata.paper_id",
                            match=models.MatchValue(value=arxiv_id),
                        )
                    ]
                )
                retriever = vector_store.as_retriever(
                    search_kwargs={"filter": filter_condition, "k": 200}
                )
                docs = await retriever._aget_relevant_documents(
                    query="*", run_manager=None
                )
                sorted_docs = await SummaryEngine._sort_docs(docs)
                full_content = "\n\n".join([doc.page_content for doc in sorted_docs])

                if not full_content.strip() and pdf_url:
                    logger.info(
                        "No indexed content found for paper %s. Falling back to PDF URL.",
                        arxiv_id,
                    )
                    full_content = await load_pdf_content(pdf_url)
            elif pdf_url:
                full_content = await load_pdf_content(pdf_url)

            if not full_content.strip():
                raise ValueError("No paper content available for summary generation")

            encoding = tiktoken.get_encoding("cl100k_base")
            tokens = encoding.encode(full_content)
            token_limit = 32768

            cumulative_summary = ""

            if len(tokens) > token_limit:
                logger.warning(
                    f"Content for paper {arxiv_id} exceeds token limit ({len(tokens)}/{token_limit}). Generating summary in chunks."
                )
                for i in range(0, len(tokens), token_limit):
                    chunk_tokens = tokens[i : i + token_limit]
                    chunk_content = encoding.decode(chunk_tokens)
                    summary = await SummaryEngine.__generate_summary(
                        chunk_content, request, model_name
                    )
                    cumulative_summary += summary + "\n\n"
                final_summary = await SummaryEngine.__generate_summary(
                    cumulative_summary, request, model_name
                )
            else:
                final_summary = await SummaryEngine.__generate_summary(
                    full_content, request, model_name
                )
            return final_summary
        except Exception as e:
            logger.error(f"Error generating summary for paper {arxiv_id}: {str(e)}")
            raise e

    @classmethod
    async def __generate_summary(
        cls, content: str, request: Optional[Request] = None,
        model_name: Optional[str] = None,
    ) -> str:
        """Generate a summary of the given content."""
        try:
            llm = LLMFactory.build_llm(
                model_name=model_name or cls.DEFAULT_MODEL,
                max_tokens=4096,
                reasoning="hidden",
                request=request,
            )
            messages = [
                {
                    "role": "system",
                    "content": PAPER_SUMMARY_SYSTEM_PROMPT,
                },
                {"role": "user", "content": content},
            ]
            response = await llm.ainvoke(messages)
            if isinstance(response.content, str):
                return response.content
            else:
                if isinstance(response.content[0], str):
                    return str(response.content[0])
                else:
                    return str(response.content[0]["text"])
        except Exception as e:
            logger.error(f"Error in __generate_summary: {str(e)}")
            raise e

    @classmethod
    async def _sort_docs(cls, docs):
        """Sort documents by page number."""

        def extract_page_number(doc):
            try:
                return int(doc.metadata.get("page_number"))
            except (ValueError, TypeError):
                return 0

        return sorted(docs, key=extract_page_number)