from typing import TypedDict

from langgraph.graph import END, StateGraph

from backend.agents.planner_agent import plan_query
from backend.services.repo_service import repo_service


class AgentState(TypedDict):
    repo_id: str
    question: str
    top_k: int
    plan: list[str]
    has_evidence: bool
    response: object


def _plan(state: AgentState) -> AgentState:
    state["plan"] = plan_query(state["question"])
    return state


def _retrieve_and_answer(state: AgentState) -> AgentState:
    from backend.services.query_service import query_service

    response = query_service.answer(state["repo_id"], state["question"], state["top_k"])
    state["has_evidence"] = bool(response.sources)
    state["response"] = response
    return state


def build_workflow():
    graph = StateGraph(AgentState)
    graph.add_node("planner", _plan)
    graph.add_node("retriever_answerer", _retrieve_and_answer)
    graph.set_entry_point("planner")
    graph.add_edge("planner", "retriever_answerer")
    graph.add_edge("retriever_answerer", END)
    return graph.compile()


workflow = build_workflow()


def run_workflow(repo_id: str, question: str, top_k: int = 8):
    from backend.services.query_service import query_service

    if not repo_service.get_chunks(repo_id):
        return query_service.answer(repo_id, question, top_k)
    state = workflow.invoke(
        {
            "repo_id": repo_id,
            "question": question,
            "top_k": top_k,
            "plan": [],
            "has_evidence": False,
            "response": None,
        }
    )
    return state["response"]