from langchain_core.documents import Document
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_core.prompts import (
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
    SystemMessagePromptTemplate,
)
from app.core.config import settings
from app.services.pinecone_service import pinecone_service
from pydantic import BaseModel, Field
from typing import AsyncGenerator, Optional
import json
import asyncio
from langchain_core.callbacks.base import BaseCallbackHandler
from langchain_core.utils.json import parse_partial_json

from app.utils.extra import process_sources


class RAGOutput(BaseModel):
    answer: str = Field(
        description=(
            "A clear, accurate, and self-contained answer to the user's question, "
            "based only on the provided context. Do not include information not present in the context."
        )
    )
    
    sources: list[int] = Field(
        default_factory=list,
        description=(
            "List of source IDs from the provided context that directly support the answer. "
            "Only include IDs that were explicitly used. Do not fabricate or infer IDs."
        )
    )
    
    video_timestamp: Optional[str] = Field(
        default=None,
        description=(
            "Timestamp in the format 'MM:SS' or 'HH:MM:SS' indicating where the answer appears in the video. "
            "Return null if no video reference is available."
        )
    )
    
    video_url: Optional[str] = Field(
        default=None,
        description=(
            "Direct URL of the video source containing the answer. "
            "Return null if no video source is available."
        )
    )

def get_prompt() -> ChatPromptTemplate:
    system_template = """You are a helpful AI assistant for Lodging System Admin .

Your primary role is to guide users step-by-step using the provided documentation context.

Instructions:
- Answer ONLY using the provided context.
- If the answer exists, explain it clearly as a sequence of steps when applicable.
- Keep responses structured, concise, and easy to follow.
- Use bullet points or numbered steps for procedures.
- If relevant, include important options, fields, or actions the user must take.
- Do NOT make up information or assumptions.
- If the answer is not found in the context:
  - Say you couldn't find it in the documentation. 
  - Ask the user to rephrase or provide more details.

Language:
- Respond in the same language as the user’s question.
- If unclear, respond in English.

Context:
{context}

Response format:
- Use markdown
- Prefer step-by-step instructions for "how-to" questions
- Highlight important actions in **bold**
"""

    messages = [
        SystemMessagePromptTemplate.from_template(system_template),
        HumanMessagePromptTemplate.from_template("{query}"),
    ]

    return ChatPromptTemplate.from_messages(messages)

class RAGService:
    def __init__(self):
        self.tokens = ""
        self.answer = ""

    def _get_embeddings(self) -> OpenAIEmbeddings:
        if not settings.OPENAI_API_KEY:
            raise ValueError("OpenAI API key is not configured.")

        return OpenAIEmbeddings(
            model=settings.EMBEDDING_MODEL,
            openai_api_key=settings.OPENAI_API_KEY,
            dimensions=512,
        )

    async def _safe_retrieve(self, query: str, k: int = 4) -> list[Document]:
        """Query Pinecone directly and tolerate vectors missing metadata/content."""
        index = pinecone_service.get_index()
        embeddings = self._get_embeddings()
        query_vector = await embeddings.aembed_query(query)

        results = index.query(
            vector=query_vector,
            top_k=k,
            include_metadata=True,
            include_values=False,
        )

        matches = getattr(results, "matches", []) or []
        documents: list[Document] = []
        for index, match in enumerate(matches):
            metadata = getattr(match, "metadata", None) or {}
            content = metadata.get("content") or metadata.get("text")
            metadata.pop("content", None)
            if not content:
                continue
            documents.append(Document(page_content=content, metadata=metadata))

        return documents

    async def _safe_retrieve_batch(self, queries: list[str], k: int = 4) -> list[Document]:
        """Batch-embed queries, then run Pinecone similarity searches concurrently."""
        cleaned_queries = [q.strip() for q in queries if isinstance(q, str) and q.strip()]
        if not cleaned_queries:
            return []

        index = pinecone_service.get_index()
        embeddings = self._get_embeddings()
        query_vectors = await embeddings.aembed_documents(cleaned_queries)

        async def _query_index(vector: list[float]):
            return await asyncio.to_thread(
                index.query,
                vector=vector,
                top_k=k,
                include_metadata=True,
                include_values=False,
            )

        results = await asyncio.gather(*(_query_index(vector) for vector in query_vectors))

        documents: list[Document] = []
        seen_keys: set[tuple[str, str]] = set()

        for result in results:
            matches = getattr(result, "matches", []) or []
            for match in matches:
                metadata = getattr(match, "metadata", None) or {}
                content = metadata.get("content") or metadata.get("text")
                metadata.pop("content", None)
                if not content:
                    continue

                key = (content, metadata.get("source") or metadata.get("url") or "")
                if key in seen_keys:
                    continue

                seen_keys.add(key)
                documents.append(Document(page_content=content, metadata=metadata))

        return documents

    def format_documents(self, docs: list[Document]) -> str:
        multiline_string = ""

        for doc in docs:
            metadata: dict = doc.metadata or {}

            title = metadata.get("title", "N/A")
            content = (doc.page_content or "").strip()
            url = metadata.get("source", None) or metadata.get("url", None)
            
            source_string = "\n<source>\n"
            title_string = f"<title>{title}</title>\n"
            content_string = f"<content>{content}</content>\n"
            url_string = f"<url>{url}</url>\n"

            multiline_string += (
                source_string
                + title_string
                + content_string
                + url_string
                + "</source>\n"
            )

        return multiline_string
    
    
    async def _process_chunk(self, token: dict) -> AsyncGenerator[str, None]:
        if token is None or token=="" or len(token)<=0:
            return
        self.tokens += token
        escaped_text = None
        
        if self.tokens:
            escaped_text = parse_partial_json(self.tokens)
            
        if escaped_text is not None and 'answer' in escaped_text:
            raw_answer = escaped_text["answer"]
            new_token = raw_answer[len(self.answer) :]
            self.answer = raw_answer
            if new_token == "" or new_token is None:
                return
            return new_token
        return None
        
    async def stream_answer(self, query: str) -> AsyncGenerator[str, None]:
        """True token streaming over SSE with final structured payload."""
        
        docs = await self._safe_retrieve(query, k=2)
        context_text = "<context>" + self.format_documents(docs) + "</context>"
        inputs = {"context": context_text, "query": query}
        callback_handler = [BaseCallbackHandler()]
        llm = ChatOpenAI(
            model=settings.LLM_MODEL,
            temperature=0.5,
            openai_api_key=settings.OPENAI_API_KEY,
            streaming=True,
            callbacks=callback_handler,
        )
        stream_prompt = get_prompt()
        token_chain = stream_prompt | llm.with_structured_output(RAGOutput)
        collected_answer = ""
        try:

            async for chunk in token_chain.astream_events(inputs):
                event_type = chunk.get("event", "")
                if event_type == "on_chain_start":
                    yield f"data: {json.dumps({'event': 'start', 'answer': 'Thinking...'})}\n\n"
                if event_type == "on_chat_model_stream":
                    chunk = chunk["data"]["chunk"]
                    if chunk.content:
                        new_token = await self._process_chunk(chunk.content)
                        if new_token is not None:
                            yield f"data: {json.dumps({'event': 'token', 'answer': new_token})}\n\n"
                elif event_type == "on_chain_end":
                   collected_answer = chunk["data"]["output"]
                   final_payload = collected_answer.model_dump()
            final_response = {
                "answer": final_payload["answer"],
                "sources": process_sources(docs, final_payload["sources"]),
                "video_timestamp": final_payload["video_timestamp"],
                "video_url": final_payload["video_url"]
            }
            yield f"data: {json.dumps({'event': 'end', 'final_response': final_response})}\n\n"
        except Exception as e:
            yield f"data: {json.dumps({'event': 'error', 'error': str(e)})}\n\n"


def get_rag_service():
    return RAGService()