from langgraph.types import interrupt
import logging
from typing import TypedDict, cast

from langchain_core.language_models import LanguageModelInput
from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import Runnable

from vdash.modules.llm import BaseLlm

from ..state import ChartState, IntentClassificationState, MetadataState, IntentEnum

system_prompt_intent_classification = """
### SYSTEM ROLE
You are an expert Intent Classifier AI specialized in chart management systems. Your primary function is to analyze user inputs and determine if they relate to customizing, formatting, or editing a chart.

### INSTRUCTION
Classify the user's input into exactly one of two categories: "EDIT_CHART" or "OTHER".

### CLASSIFICATION DEFINITIONS
1. "EDIT_CHART": The user is requesting a modification, customization, data update, or structural change to an existing chart.
   - Examples: "Change the column color to red", "Add a title to the chart", "Switch this to a line chart", "Make the font bigger".
2. "OTHER": The user's input is a general question, casual conversation, or any request entirely unrelated to chart editing.
   - Examples: "Hello", "What's the weather like today?", "Translate this sentence for me".

### LANGUAGE & TONE PROTOCOL
* **Language**: The `message_to_user` and `suggestions` MUST be in the **exact same language** as the user's input.
* **Tone**: Friendly, professional, and constructive.

### RULES & CONSTRAINTS
- Handling "OTHER": 
  1. You MUST generate a brief, polite `message_to_user` acknowledging their input while gently reminding them of your primary function (chart management).
  2. You MUST generate exactly 3 short, actionable phrases in the `suggestions` array. **These suggestions must focus STRICTLY on styling and visual formatting (e.g., changing colors, themes, font styles, or visual effects).** They must be phrased from the user's perspective (as quick-reply buttons).
- Handling "EDIT_CHART": The `message_to_user` must be an empty string ("") and the `suggestions` array must be strictly empty ([]).
- No yapping: Output strictly the JSON object. Do not include introductory text, polite phrases, or markdown formatting blocks (like ```json).

### JSON OUTPUT FORMAT
Return a single JSON object. Do not wrap it in markdown code blocks.

{{
  "intent": "EDIT_CHART" | "OTHER",
  "reason": "A brief 1-sentence explanation of why you chose this intent.",
  "message_to_user": "A polite response if OTHER, or empty string '' if EDIT_CHART.",
  "suggestions": ["Actionable style prompt 1", "Actionable style prompt 2", "Actionable style prompt 3"]
}}

### FEW-SHOT EXAMPLES

<example>
User: Can you write a poem about the sea?
Output:
{{
  "intent": "OTHER",
  "reason": "Asking for a poem is completely unrelated to chart editing.",
  "message_to_user": "I'm much better at visualizing data than writing poems! How about we focus on styling your chart to make it look great?",
  "suggestions": [
    "Change the chart colors",
    "Apply a dark theme",
    "Make the axis font bolder"
  ]
}}
</example>

<example>
User: Make the font bigger on the X axis
Output:
{{
  "intent": "EDIT_CHART",
  "reason": "The user wants to adjust the font size on an axis, which is a chart formatting task.",
  "message_to_user": "",
  "suggestions": []
}}
</example>

### CRITICAL CONSTRAINTS
1.  **NO Markdown**: Output raw JSON string only.
2.  **Valid JSON**: Ensure strictly valid JSON format.
"""

logger = logging.getLogger(__name__)


class IntentClassificationOutput(TypedDict):
    intent_classification: IntentClassificationState
    history_messages: list[AnyMessage]
    metadata: MetadataState


class HumanIntentClassificationOutput(TypedDict):
    user_input: str
    intent_classification: IntentClassificationState
    metadata: MetadataState


def route_by_intent(state: ChartState) -> IntentEnum:
    if state.intent_classification:
        return state.intent_classification.intent
    return IntentEnum.OTHER


def intent_classification_node(
    state: ChartState, llm: BaseLlm
) -> IntentClassificationOutput:
    prompt_template = ChatPromptTemplate.from_messages(
        [
            ("system", system_prompt_intent_classification),
            MessagesPlaceholder("history_messages"),
            ("human", "{user_input}"),
        ],
    )

    model_llm = cast(
        Runnable[LanguageModelInput, IntentClassificationState],
        llm.langchain.with_structured_output(IntentClassificationState),
    )
    chain = prompt_template | model_llm
    response = chain.invoke(
        {
            "user_input": state.user_input,
            "history_messages": state.history_messages,
        }
    )

    logger.info(f"Chart intent classification node: {response}")

    history_messages: list[AnyMessage] = []
    if response.intent == IntentEnum.OTHER:
        history_messages.extend(
            [
                HumanMessage(content=state.user_input),
                AIMessage(content=response.model_dump_json()),
            ]
        )

    return {
        "intent_classification": response,
        "history_messages": history_messages,
        "metadata": MetadataState(node="intent_classification_node"),
    }


def human_intent_classification_node(
    state: ChartState,
) -> HumanIntentClassificationOutput:
    user_input = interrupt(
        {
            "title": "Review question",
            "intent_classification": state.intent_classification,
        }
    )

    return {
        "user_input": user_input,
        "intent_classification": IntentClassificationState(
            intent=IntentEnum.EDIT_CHART, reason="User corrected the input"
        ),
        "metadata": MetadataState(node="human_intent_classification_node"),
    }
