from typing import Literal

from langchain_core.messages import AIMessage, SystemMessage
from langchain_core.tools import StructuredTool
from langchain_groq import ChatGroq
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import END, START, MessagesState, StateGraph
from langgraph.prebuilt import ToolNode
from pydantic import BaseModel, Field
from sqlalchemy.orm import Session

from app.config import get_settings
from app.models import User
from app.services.document_service import DocumentService
from app.services.vector_store import VectorStoreService
from app.services.web_search import build_web_search_tool


class VectorSearchInput(BaseModel):
    query: str = Field(..., description="The user question to answer from uploaded documents.")


LANGGRAPH_CHECKPOINTER = MemorySaver()


def _route_tools(state: MessagesState) -> Literal["tools", "__end__"]:
    last = state["messages"][-1]
    if isinstance(last, AIMessage) and last.tool_calls:
        return "tools"
    return "__end__"


def build_agent(*, db: Session, user: User):
    settings = get_settings()
    if not settings.groq_api_key:
        raise RuntimeError("GROQ_API_KEY is required for agent responses.")
    llm = ChatGroq(api_key=settings.groq_api_key, model=settings.model_name, temperature=0)
    document_service = DocumentService()
    vector_store = VectorStoreService()
    web_search_tool = build_web_search_tool()

    def vector_search(query: str) -> str:
        resolved_hashes = document_service.resolve_relevant_document_hashes(db, user=user, query=query)
        if not resolved_hashes:
            return "No uploaded documents are available for this user."
        matches = vector_store.similarity_search(db=db, query=query, file_hashes=resolved_hashes, k=settings.retrieval_k)
        if not matches:
            return f"No vector matches found for hashes: {resolved_hashes}"
        lines = ["Vector evidence (cite document + page + excerpt in final answer):"]
        for index, match in enumerate(matches, start=1):
            page_number = match["metadata"].get("page_number")
            page_label = str(page_number) if page_number is not None else "unknown"
            document_id = match["metadata"].get("document_id")
            score_parts = [f"distance={match['distance']:.4f}"]
            if "rerank_score" in match:
                score_parts.append(f"rerank_score={match['rerank_score']:.4f}")
            lines.append(f"{index}. document_id={document_id} | document={match['metadata']['filename']} | page={page_label} | {' | '.join(score_parts)}")
            lines.append(f"   excerpt: {match['content'][:900].replace(chr(10), ' ')}")
        return "\n\n".join(lines)

    vector_tool = StructuredTool.from_function(
        func=vector_search,
        name="vector_search",
        description=(
            "Searches the current user's uploaded documents. "
            "The tool automatically resolves the most relevant documents for the current user before chunk retrieval."
        ),
        args_schema=VectorSearchInput,
    )

    tools = [vector_tool]
    prompt = (
        "You are a document QA agent. Prefer vector_search for questions about the user's uploaded documents. "
        "Do NOT include any 'Sources' section, citation list, footnotes, chunk ids, or hashes in the final answer text. "
        "Only provide the concise user-facing answer. "
        "Citation metadata is handled separately by the application. "
        "Do not claim evidence that is not present in tool outputs."
    )
    if web_search_tool is not None:
        tools.append(web_search_tool)
        prompt += " Use web search only when the answer depends on external or current information."
    else:
        prompt += " Web search is currently unavailable in this environment."

    llm_with_tools = llm.bind_tools(tools)
    tool_node = ToolNode(tools)
    system_prompt = SystemMessage(content=prompt)

    def agent_node(state: MessagesState):
        response = llm_with_tools.invoke([system_prompt, *state["messages"]])
        return {"messages": [response]}

    graph = StateGraph(MessagesState)
    graph.add_node("agent", agent_node)
    graph.add_node("tools", tool_node)
    graph.add_edge(START, "agent")
    graph.add_conditional_edges("agent", _route_tools, {"tools": "tools", "__end__": END})
    graph.add_edge("tools", "agent")

    return graph.compile(checkpointer=LANGGRAPH_CHECKPOINTER)