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 langgraph.types import interrupt

from vdash.modules.llm import BaseLlm

from ..state import (
    IntentClassificationState,
    IntentEnum,
    MetadataState,
    SqlState,
)

system_prompt_intent_classification = """
### SYSTEM ROLE
You are an expert Data Analyst and SQL Intent Classifier. Your goal is to accurately categorize user natural language queries into specific SQL-related intents.

### DATABASE SCHEMA
{database_schemas}

### INSTRUCTION
Analyze the user's input (provided in the user message) and follow these steps:
1.  **Language Detection**: Identify the language of the user's input.
2.  **Entity Analysis**: Analyze the input for specific SQL entities (metrics, dimensions, timeframes, filters).
3.  **Classification**: Determine if the input is `DIRECT_SQL`, `AMBIGUOUS_SQL`, or `IRRELEVANT`.
4.  **Output Generation**: Construct a valid JSON object following the strict schema rules below.

### CLASSIFICATION DEFINITIONS
1.  **DIRECT_SQL**: The query contains sufficient entities to construct a WHERE clause and SELECT statement.
2.  **AMBIGUOUS_SQL**: The query is related to business data but lacks specific constraints.
3.  **IRRELEVANT**: Social chitchat, greetings, or questions unrelated to the database schema.

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

### OUTPUT SCHEMA RULES (CRITICAL)
1.  **If DIRECT_SQL**:
    * `message_to_user`: null
    * `clarification`: null
2.  **If IRRELEVANT**:
    * `message_to_user`: A friendly conversational response.
    * `clarification`: null
3.  **If AMBIGUOUS_SQL**:
    * `message_to_user`: A polite question asking for missing info.
    * `clarification`: An object containing `ambiguity_reason` and `suggestions`.

### JSON OUTPUT FORMAT
Return a single JSON object. Do not wrap it in markdown code blocks.
Note: The "clarification" field must be null if intent is NOT AMBIGUOUS_SQL.
{{
  "intent": "DIRECT_SQL | AMBIGUOUS_SQL | IRRELEVANT",
  "message_to_user": "String or null (based on rules above)",
  "analysis": "Brief reasoning for the classification",
  "clarification": {{
    "ambiguity_reason": "Explanation of missing entities",
    "suggestions": ["Option 1", "Option 2", "Option 3"]
  }}
}}

### FEW-SHOT EXAMPLES

<example>
User: Show me revenue for the top 5 products last month.
Output:
{{
  "intent": "DIRECT_SQL",
  "message_to_user": null,
  "analysis": "Input contains specific metric (revenue), limit (top 5), dimension (products), and timeframe (last month).",
  "clarification": null
}}
</example>

<example>
User: Tình hình kinh doanh thế nào?
Output:
{{
  "intent": "AMBIGUOUS_SQL",
  "message_to_user": "Để tôi giúp bạn xem báo cáo nhé. Tuy nhiên, bạn muốn xem cụ thể về Doanh thu, Lợi nhuận hay Số lượng đơn hàng, và trong khoảng thời gian nào?",
  "analysis": "Người dùng hỏi chung chung, thiếu chỉ số và thời gian.",
  "clarification": {{
    "ambiguity_reason": "Thiếu chỉ số (Metric) và Thời gian (Timeframe).",
    "suggestions": [
      "Xem doanh thu tháng này",
      "Xem lợi nhuận theo quý",
      "Xem số đơn hàng hôm nay"
    ]
  }}
}}
</example>

<example>
User: Chào em, nay mệt không?
Output:
{{
  "intent": "IRRELEVANT",
  "message_to_user": "Cảm ơn bạn đã hỏi thăm, tôi là AI nên luôn sẵn sàng hỗ trợ bạn! Bạn có cần xem dữ liệu gì hôm nay không?",
  "analysis": "User is greeting and asking about status, irrelevant to data.",
  "clarification": null
}}
</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: SqlState) -> IntentEnum:
    if state.intent_classification:
        return state.intent_classification.intent
    return IntentEnum.IRRELEVANT


def intent_classification_node(
    state: SqlState, 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,
            "database_schemas": state.database_schemas,
            "history_messages": state.history_messages,
        }
    )

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

    history_messages: list[AnyMessage] = []
    if response.intent != IntentEnum.DIRECT_SQL:
        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: SqlState,
) -> HumanIntentClassificationOutput:
    user_input = interrupt(
        {
            "title": "Review question",
            "intent_classification": state.intent_classification,
        }
    )

    return {
        "user_input": user_input,
        "intent_classification": IntentClassificationState(
            intent=IntentEnum.DIRECT_SQL,
            message_to_user=None,
            analysis="User corrected the input",
            clarification=None,
        ),
        "metadata": MetadataState(node="human_intent_classification_node"),
    }
