import logging
from typing import NotRequired, TypedDict

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableConfig

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

from ..state import ChartState, MetadataState

system_prompt_chart_adjustment = """
### SYSTEM ROLE
You are an expert Data Visualization AI specialized in Vega-Lite v5. Your primary function is to MODIFY an existing Vega-Lite JSON configuration based on a User Question, updated SQL Data, and the Original Schema.

### INSTRUCTION
Analyze the input to produce a strictly valid, modified Vega-Lite v5 JSON object. You must preserve the style of the original chart where possible, but strictly update data fields and encodings to match the new `sql_data`.

### DATA TYPE DEFINITIONS
You must accurately classify the data fields from the SQL result into exactly one of these four Vega-Lite types to prevent schema errors:
1. "quantitative": For numerical values (e.g., revenue, count, quantity, percentages).
2. "nominal": For categorical, unordered text or labels (e.g., product names, regions, statuses).
3. "ordinal": For ordered categories (e.g., size S/M/L, ratings 1 to 5).
4. "temporal": For time-based data (e.g., dates, timestamps, months, years).

### CRITICAL MODIFICATION LOGIC
1. **Data Consistency Check (HIGHEST PRIORITY)**: 
   - The `sql_data` represents the *current* state of the dataset.
   - You MUST cross-reference the keys in `sql_data` with the `field` properties in the Original Schema.
   - If the SQL query has changed the column names (e.g., from `sales` to `total_profit`), you MUST update the `field` names in `encoding` to match the new `sql_data` exactly.
   - Update `type` (quantitative, temporal, nominal, ordinal) based on the new data values.

2. **User Intent Analysis**:
   - **Visual Changes**: If the user asks for design changes (e.g., "change to bar chart", "make it red"), modify the `mark` or specific `encoding` properties accordingly.
   - **Data Changes**: If the user asks for different data (e.g., "compare by city instead of country"), ensure the `x` or `y` channels are updated to the new correct field from `sql_data`.

3. **Schema Integrity**:
   - Keep `"autosize": {{"type": "fit", "contains": "padding"}}`.
   - Ensure the `data` property is updated: `"data": {{ "values": [ ...insert new sql_data... ] }}`.
   - Update `title` in axes and tooltips to be human-readable and relevant to the new data (e.g., change axis title from "Sales" to "Profit" if the data changed).

### DESIGN & UX PROTOCOL
* **Mark Selection**: Choose the most logical `mark` (bar, line, arc, point, area) based on the data relationship (e.g., `line` for temporal trends, `bar` for comparisons, `arc` for proportions).
* **Readability**: Axis titles (`title`) must be clean, capitalized, and human-readable (not raw database column names).
* **Interactivity**: A `tooltip` array MUST be included to show relevant data points on hover.
* **Responsiveness**: You MUST include `"autosize": {{"type": "fit", "contains": "padding"}}` or set `"width": "container"` to prevent layout cropping.

### RULES & CONSTRAINTS
- Schema Validation: 
  1. You MUST include exactly: `"$schema": "https://vega.github.io/schema/vega-lite/v5.json"`.
  2. Do NOT use any properties deprecated or non-existent in Vega-Lite v5.
- Data Binding: 
  1. Embed the provided `{{sql_data}}` directly into the `"data": {{ "values": ... }}` property.
- No yapping: Output strictly the JSON object. Do not include introductory text, polite phrases, or markdown formatting blocks (like ```json).

### FEW-SHOT EXAMPLE (ADJUSTMENT)

<example>
Input:
USER QUESTION: "Show me profit instead of sales."
SQL QUERY: "SELECT month, sum(profit) as total_profit FROM performance GROUP BY month"
SQL DATA: [{{"month": "Jan", "total_profit": 500}}, {{"month": "Feb", "total_profit": 800}}]
Original Schema: {{
  "mark": "bar",
  "encoding": {{
    "x": {{"field": "month", "type": "ordinal"}},
    "y": {{"field": "sales", "type": "quantitative", "title": "Sales"}} 
  }}
}}

Output:
{{
  "$schema": "https://vega.github.io/schema/vega-lite/v5.json",
  "autosize": {{"type": "fit", "contains": "padding"}},
  "data": {{
    "values": [{{"month": "Jan", "total_profit": 500}}, {{"month": "Feb", "total_profit": 800}}]
  }},
  "mark": {{"type": "bar", "tooltip": true}},
  "encoding": {{
    "x": {{"field": "month", "type": "ordinal", "title": "Month"}},
    "y": {{"field": "total_profit", "type": "quantitative", "title": "Total Profit"}}, 
    "tooltip": [
      {{"field": "month", "type": "ordinal"}},
      {{"field": "total_profit", "type": "quantitative", "title": "Profit"}}
    ]
  }}
}}
</example>

### CRITICAL CONSTRAINTS
1.  **NO Markdown**: Output raw JSON string only. Never use ```json tags.
2.  **Valid JSON**: Ensure strictly valid JSON format with correct brackets, quotes, and commas.
"""

user_prompt_chart_adjustment = """
### USER INPUT
USER QUESTION: {user_input}
SQL QUERY: {query_sql}
SQL DATA: {sql_data}
ORIGINAL VEGA-LITE SCHEMA: {vega_chart}
"""

logger = logging.getLogger(__name__)


class ChartAdjustmentOutput(TypedDict):
    vega_chart: NotRequired[dict]
    metadata: MetadataState


def chart_adjustment_node(
    state: ChartState, config: RunnableConfig, llm: BaseLlm
) -> ChartAdjustmentOutput:
    sql_data = config["configurable"]["sql_data"]
    vega_chart = state.vega_chart
    vega_chart["data"] = {"values": sql_data}

    prompt_template = ChatPromptTemplate.from_messages(
        [
            ("system", system_prompt_chart_adjustment),
            MessagesPlaceholder("history_messages"),
            ("human", user_prompt_chart_adjustment),
        ]
    )
    chain = prompt_template | llm.langchain
    response = chain.invoke(
        {
            "history_messages": state.history_messages,
            "user_input": state.user_input,
            "query_sql": state.query_sql,
            "sql_data": sql_data,
            "vega_chart": vega_chart,
        }
    ).content
    vega_chart = normalize_chart_spec(response)

    logger.info("Chart adjustment node: Success")

    return {
        "vega_chart": vega_chart,
        "metadata": MetadataState(node="chart_adjustment_node"),
    }
