from functools import partial
from typing import Any, Optional, cast

from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.graph import END, START, StateGraph
from langgraph.types import Command

from vdash.modules.llm import BaseLlm
from vdash.utils import SqlClient

from .nodes import (
    human_intent_classification_node,
    intent_classification_node,
    route_by_intent,
    sql_correction_node,
    sql_generation_node,
)
from .state import IntentEnum, SqlState


class SqlAgent:
    def __init__(
        self,
        llm: Optional[BaseLlm] = None,
        checkpointer: Optional[BaseCheckpointSaver] = None,
    ):
        if llm is None:
            raise ValueError("llm must be provided")
        self.llm = llm

        if checkpointer is None:
            raise ValueError("checkpointer must be provided")
        self.checkpointer = checkpointer
        self.app = self.build_workflow()

    def build_workflow(self):
        workflow = StateGraph(state_schema=SqlState)

        workflow.add_node(
            "intent_classification_node",
            partial(intent_classification_node, llm=self.llm),
        )
        workflow.add_node(
            "sql_generation_node",
            partial(sql_generation_node, llm=self.llm),
        )
        workflow.add_node(
            "sql_correction_node",
            partial(sql_correction_node, llm=self.llm),
        )
        workflow.add_node(
            "human_intent_classification_node", human_intent_classification_node
        )

        workflow.add_edge(START, "intent_classification_node")
        workflow.add_conditional_edges(
            "intent_classification_node",
            route_by_intent,
            {
                IntentEnum.DIRECT_SQL: "sql_generation_node",
                IntentEnum.AMBIGUOUS_SQL: "human_intent_classification_node",
                IntentEnum.IRRELEVANT: END,
            },
        )
        workflow.add_edge("human_intent_classification_node", "sql_generation_node")
        workflow.add_edge("sql_generation_node", "sql_correction_node")
        workflow.add_edge("sql_correction_node", END)

        return workflow.compile(checkpointer=self.checkpointer)

    def ask_question(self, db_url: str, thread_id: str, user_input: str) -> SqlState:
        sql_client = SqlClient(db_url)

        initial_state = SqlState(
            dialect=sql_client.get_dialect(),
            database_schemas=sql_client.get_schema_database(),
            user_input=user_input,
        )
        config = self._get_config(thread_id, sql_client)

        result = self.app.invoke(initial_state, config)
        return SqlState.model_validate(result)

    def confirm_question(
        self, db_url: str, thread_id: str, user_input: str
    ) -> SqlState:
        sql_client = SqlClient(db_url)
        config = self._get_config(thread_id, sql_client)

        result = self.app.invoke(Command(resume=user_input), config)
        return SqlState.model_validate(result)

    def is_waiting_for_user(self, thread_id: str) -> bool:
        config = self._get_config(thread_id)
        state = self.app.get_state(config)
        return len(state.next) > 0

    def _get_config(
        self, thread_id: str, sql_client: Optional[SqlClient] = None
    ) -> RunnableConfig:
        configurable: dict[str, Any] = {"thread_id": thread_id}
        if sql_client is not None:
            configurable["sql_client"] = sql_client
        return cast(RunnableConfig, {"configurable": configurable})
