import json
import re
from typing import AsyncGenerator, Annotated, TypedDict, Sequence, Literal, Any
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage, AIMessage
from langgraph.graph.message import add_messages
from langgraph.graph import StateGraph, START, END
from langgraph.prebuilt import ToolNode
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langchain_core.utils.json import parse_partial_json
from app.core.config import settings
from app.core.prompts import AGENT_SYSTEM_PROMPT
from app.core.constants import MAX_TOKEN_LIMIT, RETRIEVE_DOCUMENTATION_DESCRIPTION
from app.services.rag_service import RAGService
from app.models.chat import AgentRetrieverResponse, ClassifyResponse


@tool(
    description=RETRIEVE_DOCUMENTATION_DESCRIPTION,
    args_schema=AgentRetrieverResponse
)
async def retrieve_documentation(query: str, question: str):
    all_queries = [query, question]
    rag_service = RAGService()
    documents = await rag_service._safe_retrieve_batch(all_queries, k=5)
    return rag_service.format_documents(documents)


# 2. Define the State for LangGraph
class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]

class LangGraphAgentService:
    MAX_TOKEN_LIMIT = MAX_TOKEN_LIMIT

    def __init__(self, tools: list | None = None):
        self.tokens = ""
        self.answer = ""
        self.agent_structured_output = {}
        # Initialize the LLM with streaming enabled
        self.llm = ChatOpenAI(
            model=settings.LLM_MODEL,
            temperature=0.5,
            openai_api_key=settings.OPENAI_API_KEY,
            streaming=True,
        )
        
        # Bind the tools to the LLM
        self.tools = [retrieve_documentation]
        self.llm_with_tools = self.llm.bind_tools(
            tools=self.tools,
            parallel_tool_calls=True,
            strict=True,
            response_format=ClassifyResponse
        )
        
        # Compile the graph
        self.graph = self._build_graph()

    def _num_tokens_from_messages(self, messages: list[BaseMessage]) -> int:
        """Approximate token count for history trimming."""
        try:
            return self.llm.get_num_tokens_from_messages(messages)
        except Exception:
            # Fallback approximation if tokenizer/model metadata is unavailable.
            total_chars = 0
            for msg in messages:
                content = msg.content
                if isinstance(content, str):
                    total_chars += len(content)
                elif isinstance(content, list):
                    for item in content:
                        if isinstance(item, dict):
                            text = item.get("text")
                            if isinstance(text, str):
                                total_chars += len(text)
            return max(1, total_chars // 4)

    def _history_to_base_messages(self, history: list[dict], question_token: int = 0) -> list[BaseMessage]:
        if not history:
            return []

        normalized: list[BaseMessage] = []
        for item in history:
            if not isinstance(item, dict):
                continue
            role = str(item.get("type", "")).lower()
            text = item.get("text", "")
            if not isinstance(text, str) or not text.strip():
                continue

            if role == "human":
                normalized.append(HumanMessage(content=text))
            elif role in {"bot", "ai", "assistant"}:
                normalized.append(AIMessage(content=text))

        if not normalized:
            return []

        token_count = self._num_tokens_from_messages(normalized)
        max_allowed = max(0, self.MAX_TOKEN_LIMIT - question_token)
        while normalized and token_count > max_allowed:
            normalized.pop(0)
            token_count = self._num_tokens_from_messages(normalized)

        return normalized
    
    def get_agent_prompt(self) -> list[BaseMessage]:
        """Creates the base prompt for the agent."""
        system_template = AGENT_SYSTEM_PROMPT
        return [SystemMessage(content=system_template)]

    async def call_agent(self, state: AgentState):
        """The agent node execution logic."""
        # Start with the system prompt
        messages = self.get_agent_prompt()
        # Append all conversation history and tool outputs
        messages.extend(state["messages"])
        
        # Invoke the LLM
        response = await self.llm_with_tools.ainvoke(
            messages,
            config={"run_name": "Agent"}
        )
        return {"messages": [response]}

    def should_continue(self, state: AgentState) -> Literal["tools", END]:
        """Conditional routing: continue to tools if the LLM called a tool, otherwise END."""
        messages = state["messages"]
        last_message = messages[-1]
        
        # If the LLM makes a tool call, route to the "tools" node
        if last_message.tool_calls:
            return "tools"
        # Otherwise, end the graph execution
        return END

    def _build_graph(self):
        """Construct the state graph."""
        workflow = StateGraph(AgentState)
        
        # Add the agent and tools nodes
        workflow.add_node("agent", self.call_agent)
        workflow.add_node("tools", ToolNode(self.tools))
        
        # Define the edges
        workflow.add_edge(START, "agent")
        workflow.add_conditional_edges("agent", self.should_continue)
        workflow.add_edge("tools", "agent")
        
        # Compile into a runnable LangGraph application
        return workflow.compile()
    
    def _normalize_chunk_content(self, content) -> str:
        """Normalize LangChain chunk content to plain text."""
        if content is None:
            return ""
        if isinstance(content, str):
            return content
        if isinstance(content, list):
            parts = []
            for item in content:
                if isinstance(item, str):
                    parts.append(item)
                elif isinstance(item, dict):
                    text = item.get("text")
                    if isinstance(text, str):
                        parts.append(text)
            return "".join(parts)
        if isinstance(content, dict):
            text = content.get("text")
            return text if isinstance(text, str) else ""
        return ""

    def _strip_source_markup(self, text: str) -> str:
        """Remove source tags/markers from streamed answer text."""
        if not text:
            return ""
        # Remove XML-like source tags such as <source id=0>...</source>
        cleaned = re.sub(r"</?source\b[^>]*>", "", text, flags=re.IGNORECASE)
        # Remove escaped variants that can appear in streamed chunks.
        cleaned = cleaned.replace("\\n<source id=0>\\n", "\\n")
        cleaned = cleaned.replace("\n<source id=0>\n", "\n")
        return cleaned

    def _has_source_marker(self, text: str) -> bool:
        """Check whether token text contains source marker syntax."""
        if not text:
            return False
        lowered = text.lower()
        return "<source" in lowered or "\\n<source" in lowered

    async def _process_chunk(self, token: str) -> str:
        if not token:
            return None

        # If the model streams JSON (response_format), buffer until parsable.
        stripped = token.lstrip()
        if not self.tokens:
            if stripped.startswith("{"):
                self.tokens = token
                return None
            # Non-JSON text stream: pass through directly.
            return token
        else:
            self.tokens += token

        try:
            escaped_text = parse_partial_json(self.tokens)
        except Exception:
            # Ignore incomplete JSON fragments like "{", "\"ans", etc.
            return None

        if not isinstance(escaped_text, dict):
            return None

        self.agent_structured_output = escaped_text
        raw_answer = escaped_text.get("answer")
        if raw_answer is not None:
            if not isinstance(raw_answer, str):
                return None
            clean_answer = self._strip_source_markup(raw_answer)
            new_token = clean_answer[len(self.answer):] if clean_answer.startswith(self.answer) else clean_answer
            self.answer = clean_answer
            if not new_token:
                return None
            return new_token
        return None

    def _parse_stream_event(self, event: Any) -> tuple[str, dict]:
        """Convert LangGraph stream events into token payloads."""
        if not isinstance(event, tuple) or len(event) != 2:
            return "", {}

        mode, payload = event
        if mode != "messages":
            return "", {}
        if not isinstance(payload, tuple) or len(payload) != 2:
            return "", {}

        chunk, meta = payload
        chunk_content = self._normalize_chunk_content(getattr(chunk, "content", ""))
        token_text = chunk_content if isinstance(chunk_content, str) else ""
        return token_text

    def _get_video_url_for_language(self, video_url: str, language: str) -> str:
        if not video_url:
            return None

        lang = (language or "en").lower()
        if lang not in {"en", "hi", "gu"}:
            lang = "en"

        if lang == "hi":
            lang = "gu" # Map Hindi to Gujarati folder for video URLs

        # Replace language folder in URL
        return video_url.replace("/en/", f"/{lang}/")
    
    
    async def stream_answer(self, query: str, history: list[dict] | None = None) -> AsyncGenerator[str, None]:
        self.tokens = ""
        self.answer = ""
        self.agent_structured_output = {}
        question_token = self._num_tokens_from_messages([HumanMessage(content=query)])
        history_messages = self._history_to_base_messages(history or [], question_token=question_token)
            
        inputs = {"messages": history_messages}
        
        yield f"data: {json.dumps({'event': 'start', 'answer': 'Thinking...'})}\n\n"

        try:
            # version="v2" is required for LangChain v0.2/v0.3 standard stream events
            async for event in self.graph.astream(inputs, stream_mode=["messages"]):
                token_text = self._parse_stream_event(event)

                if token_text:
                    processed = await self._process_chunk(token_text)
                    if processed:
                        if self._has_source_marker(processed):
                            continue
                        clean_token = self._strip_source_markup(processed)
                        if clean_token:
                            yield f"data: {json.dumps({'event': 'token','answer': clean_token})}\n\n"
                        
            final_response = {
                "answer": self.answer,
                "video_start_timestamp": self.agent_structured_output.get("video_start_timestamp", None),
                "video_end_timestamp": self.agent_structured_output.get("video_end_timestamp", None),
                "language": self.agent_structured_output.get("language", "en"),
                "next_possible_questions": self.agent_structured_output.get("next_possible_questions", []),
                "video_url": self.agent_structured_output.get("video_url", None),
            }
            final_response["video_url"] = self._get_video_url_for_language( self.agent_structured_output.get("video_url", None), final_response["language"])
            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_langgraph_service():
    return LangGraphAgentService()