import logging
from typing import NotRequired, TypedDict

from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableConfig

from vdash.modules.llm.base import BaseLlm
from vdash.utils import validate_chart
from vdash.utils.chart import normalize_chart_spec

from ..state import ChartState, MetadataState, MetadataStatus

system_prompt_chart_validate = """
### SYSTEM ROLE
You are an expert in evaluating and fixing Vega-Lite specifications (JSON).
Your task is to analyze the provided Vega-Lite JSON and the error message, then output a corrected, valid JSON specification.

### INPUT DATA
Vega-Lite Spec: {chart}
Error Context: {error}

### RULES
1. Output ONLY the single, final valid Vega-Lite JSON string.
2. NO Markdown formatting (do not use ```json or ```vega-lite).
3. NO comments, NO explanations, NO conversational filler.
4. Ensure the JSON syntax is valid (correct brackets, commas, and quotes).
5. Fix specific Vega-Lite errors (e.g., invalid 'mark' types, missing 'data' sources, wrong 'encoding' channels, or deprecated keywords).
6. Ensure the 'data' field is correctly formatted (either inline values or a valid URL).

### OUTPUT FORMAT
Return strictly the raw fixed JSON text.

### CRITICAL CONSTRAINTS
1. **NO MARKDOWN**: Do NOT wrap the output in markdown blocks. Return raw JSON text only.
2. **NO CONVERSATION**: Do NOT include text like "Here is the fixed chart".
3. **VALID JSON**: The output must be parseable by standard JSON parsers.
"""

logger = logging.getLogger(__name__)


class ChartValidationOutput(TypedDict):
    vega_chart: NotRequired[dict]
    metadata: MetadataState
    history_messages: NotRequired[list[AnyMessage]]


def chart_validation_node(
    state: ChartState, config: RunnableConfig, llm: BaseLlm
) -> ChartValidationOutput:
    sql_data = config["configurable"]["sql_data"]
    vega_chart = state.vega_chart
    last_error = None

    prompt_template = ChatPromptTemplate.from_messages(
        ("system", system_prompt_chart_validate)
    )
    chain = prompt_template | llm.langchain

    for _ in range(3):
        vega_chart["data"] = {"values": sql_data}
        ok, err = validate_chart(vega_chart)
        last_error = err
        if ok:
            logger.info("Validate chart node: Chart passed")

            return {
                "vega_chart": normalize_chart_spec(vega_chart),
                "metadata": MetadataState(node="chart_validation_node"),
                "history_messages": [
                    HumanMessage(content=state.user_input),
                    AIMessage(
                        content=f"Vega-Lite Schema: {normalize_chart_spec(vega_chart)}"
                    ),
                ],
            }

        logger.info(f"Validate chart node: Chart failed: {err}")
        response = chain.invoke({"chart": vega_chart, "error": err}).content
        vega_chart = normalize_chart_spec(response)

    return {
        "vega_chart": normalize_chart_spec(vega_chart),
        "metadata": MetadataState(
            node="chart_validation_node",
            status=MetadataStatus.ERROR,
            error=last_error,
        ),
    }
