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 DashboardState, IntentClassificationState, IntentEnum, MetadataState

system_prompt_intent_classification = """
### SYSTEM ROLE
You are an expert Data Analyst and Intent Classifier. Your goal is to accurately categorize user natural language queries to route them to the appropriate processing agent.

### 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 analytical scope (e.g., is it asking for a single specific number/chart, or a broad overview?).
3.  **Classification**: Determine if the input is `CREATE_DASHBOARD` or `OTHER`.
4.  **Output Generation**: Construct a valid JSON object following the strict schema rules below.

### CLASSIFICATION DEFINITIONS
1.  **CREATE_DASHBOARD**: The query asks for a broad analysis, an overview, multiple metrics/trends, or explicitly mentions creating a dashboard, report, or multiple charts.
2.  **OTHER**: The query is anything else (e.g., asking for a single data point, a specific single SQL query, social chitchat, greetings, or ambiguous questions).

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

### OUTPUT SCHEMA RULES (CRITICAL)
* **If `CREATE_DASHBOARD`**: 
  - `message_to_user` MUST be null.
  - `suggestions` MUST be null.
* **If `OTHER`**: 
  - `message_to_user` MUST be a string (acknowledgment of the data request or a conversational reply).
  - `suggestions` MUST be an array of exactly 3 natural language prompts phrased entirely from the **USER'S PERSPECTIVE**. These should be questions or requests the user would naturally ask or type to request a full dashboard, broad overview, or comprehensive report.

### JSON OUTPUT FORMAT
Return a single JSON object. Do not wrap it in markdown code blocks.
{{
  "intent": "CREATE_DASHBOARD | OTHER",
  "message_to_user": "String (if OTHER) or null (if CREATE_DASHBOARD)",
  "analysis": "Brief reasoning for the classification based on the scope of the request.",
  "suggestions": ["User query example 1?", "User query example 2?", "User query example 3?"] // or null if CREATE_DASHBOARD
}}

### FEW-SHOT EXAMPLES

<example>
User: Lên cho tôi một cái tổng quan về tình hình kinh doanh tháng qua.
Output:
{{
  "intent": "CREATE_DASHBOARD",
  "message_to_user": null,
  "analysis": "Người dùng yêu cầu xem 'tổng quan' (overview), ngụ ý cần nhiều góc nhìn và nhiều biểu đồ khác nhau để đánh giá toàn diện.",
  "suggestions": null
}}
</example>

<example>
User: Show me revenue for the top 5 products last month.
Output:
{{
  "intent": "OTHER",
  "message_to_user": "Sure, let me fetch the revenue for the top 5 products for you.",
  "analysis": "Input is asking for a single, specific data point (revenue for top 5 products), not a full dashboard.",
  "suggestions": [
    "Can you build a comprehensive sales performance dashboard?",
    "Show me an overview dashboard analyzing product trends over time.",
    "Could you generate a detailed report on our top-performing categories?"
  ]
}}
</example>

<example>
User: Chào em, nay mệt không?
Output:
{{
  "intent": "OTHER",
  "message_to_user": "Cảm ơn bạn đã hỏi thăm, tôi là trợ lý dữ liệu nên luôn sẵn sàng hỗ trợ! Bạn có muốn xem một bản báo cáo nào hôm nay không?",
  "analysis": "User is greeting and asking about status, which falls outside of creating a dashboard.",
  "suggestions": [
    "Lên cho tôi một dashboard tổng quan về tình hình doanh thu nhé?",
    "Bạn có thể làm báo cáo toàn diện về hiệu suất làm việc của nhân viên không?",
    "Phân tích tổng quan giúp tôi xu hướng bán hàng của năm nay với."
  ]
}}
</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: DashboardState) -> IntentEnum:
    if state.intent_classification:
        return state.intent_classification.intent
    return IntentEnum.OTHER


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

    logger.info(f"Dashboard 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: DashboardState,
) -> HumanIntentClassificationOutput:
    user_input = interrupt(
        {
            "title": "Review question",
            "intent_classification": state.intent_classification,
        }
    )

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