from typing import Any
from functools import partial
from typing import 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 (
    chart_adjustment_node,
    chart_generation_node,
    chart_validation_node,
    human_intent_classification_node,
    intent_classification_node,
    route_by_intent,
)
from .state import ChartState, IntentEnum, MetadataState, MetadataStatus


class ChartAgent:
    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.edit_chart_app = self.build_edit_chart_workflow()
        self.generate_chart_app = self.build_generate_chart_workflow()

    def build_edit_chart_workflow(self):
        workflow = StateGraph(state_schema=ChartState)

        workflow.add_node(
            "intent_classification_node",
            partial(intent_classification_node, llm=self.llm),
        )
        workflow.add_node(
            "human_intent_classification_node", human_intent_classification_node
        )
        workflow.add_node(
            "chart_adjustment_node",
            partial(chart_adjustment_node, llm=self.llm),
        )
        workflow.add_node(
            "chart_validation_node",
            partial(chart_validation_node, llm=self.llm),
        )

        workflow.add_edge(START, "intent_classification_node")
        workflow.add_conditional_edges(
            "intent_classification_node",
            route_by_intent,
            {
                IntentEnum.EDIT_CHART: "chart_adjustment_node",
                IntentEnum.OTHER: "human_intent_classification_node",
            },
        )
        workflow.add_edge("human_intent_classification_node", "chart_adjustment_node")
        workflow.add_edge("chart_adjustment_node", "chart_validation_node")
        workflow.add_edge("chart_validation_node", END)

        return workflow.compile(
            checkpointer=self.checkpointer,
        )

    def build_generate_chart_workflow(self):
        workflow = StateGraph(state_schema=ChartState)

        workflow.add_node(
            "chart_generation_node",
            partial(chart_generation_node, llm=self.llm),
        )
        workflow.add_node(
            "chart_validation_node",
            partial(chart_validation_node, llm=self.llm),
        )

        workflow.add_edge(START, "chart_generation_node")
        workflow.add_edge("chart_generation_node", "chart_validation_node")
        workflow.add_edge("chart_validation_node", END)

        return workflow.compile()

    def edit_chart(
        self,
        db_url: str,
        thread_id: str,
        user_input: str,
        query_sql: str,
        vega_chart: dict,
    ) -> ChartState:
        sql_client = SqlClient(db_url)
        sql_data, err = sql_client.run(query_sql)

        if not sql_data:
            return ChartState(
                user_input=user_input,
                query_sql=query_sql,
                metadata=MetadataState(
                    node="get_sql_data",
                    status=MetadataStatus.ERROR,
                    error=err,
                ),
            )

        initial_state = ChartState(
            user_input=user_input,
            query_sql=query_sql,
            vega_chart=vega_chart,
        )
        config = self._get_config(sql_data, thread_id)

        result = self.edit_chart_app.invoke(initial_state, config)
        return ChartState.model_validate(result)

    def confirm_edit_chart(
        self,
        db_url: str,
        thread_id: str,
        user_input: str,
        query_sql: str,
    ) -> ChartState:
        sql_client = SqlClient(db_url)
        sql_data, err = sql_client.run(query_sql)

        if not sql_data:
            return ChartState(
                user_input=user_input,
                query_sql=query_sql,
                metadata=MetadataState(
                    node="get_sql_data",
                    status=MetadataStatus.ERROR,
                    error=err,
                ),
            )
        config = self._get_config(sql_data, thread_id)

        result = self.edit_chart_app.invoke(Command(resume=user_input), config)
        return ChartState.model_validate(result)

    def generate_chart(
        self, db_url: str, user_input: str, query_sql: str
    ) -> ChartState:
        sql_client = SqlClient(db_url)
        sql_data, err = sql_client.run(query_sql)

        if not sql_data:
            return ChartState(
                user_input=user_input,
                query_sql=query_sql,
                metadata=MetadataState(
                    node="get_sql_data",
                    status=MetadataStatus.ERROR,
                    error=err,
                ),
            )

        initial_state = ChartState(
            user_input=user_input,
            query_sql=query_sql,
        )
        config = self._get_config(sql_data)

        return cast(ChartState, self.generate_chart_app.invoke(initial_state, config))

    def _get_config(
        self, sql_data: list[dict], thread_id: Optional[str] = None
    ) -> RunnableConfig:
        configurable: dict[str, Any] = {"sql_data": sql_data}

        if thread_id:
            configurable["thread_id"] = thread_id
        return cast(RunnableConfig, {"configurable": configurable})
