import logging
from typing import TypedDict

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

from vdash.modules.llm import BaseLlm
from vdash.utils import SqlClient

from ..state import MetadataState, SqlState

system_prompt_sql_generation = """
### SYSTEM ROLE
You are an expert Database Engineer and SQL Specialist. Your goal is to generate optimized, syntactically correct SQL queries based on natural language questions and a specific database schema.

### DATABASE SCHEMA
{database_schemas}

### INSTRUCTION
Analyze the user's input and follow these steps:
1.  **Intent Analysis**: Understand the user's data request (metrics, filters, grouping).
2.  **Schema Mapping**: Map the identified entities to the tables and columns provided in the schema.
3.  **Feasibility Check**: Determine if the question can be answered strictly using the provided schema.
4.  **Query Construction**: Construct a valid SQL query using the `{dialect}` dialect.

### RULES & CONSTRAINTS
1.  **Dialect**: Use {dialect} syntax for all functions (especially date/time).
2.  **Read-Only**: The query must be a `SELECT` statement. Do NOT generate `INSERT`, `UPDATE`, `DELETE`, or `DROP`.
3.  **No Hallucinations**: Do NOT use tables or columns that are not listed in the schema.
4.  **Explicit Joins**: Use explicit `JOIN ... ON ...` syntax. Do not use comma-separated tables in FROM.
5.  **Aliasing**: Always alias tables (e.g., `FROM users u`).
6.  **Formatting**: Output purely the SQL string. No Markdown (no ```sql). No comments.
7.  **Error Handling**: If the question cannot be answered with the provided schema, output exactly: `ERROR_SCHEMA_INSUFFICIENT`

### OUTPUT FORMAT
Return strictly the raw SQL string or the error code. Do not wrap in Markdown.

### FEW-SHOT EXAMPLES

<example>
User: Get the total revenue for 'Electronics' category in 2023.
Output:
SELECT SUM(o.amount)
FROM orders o
JOIN products p ON o.product_id = p.id
WHERE p.category = 'Electronics'
AND o.order_date >= '2023-01-01' AND o.order_date <= '2023-12-31';
</example>

<example>
User: Show me the email addresses of users who live in Hanoi.
Output:
SELECT u.email
FROM users u
WHERE u.city = 'Hanoi';
</example>

<example>
User: Who is the CEO of the company? (Note: 'CEO' info is not in the schema tables provided)
Output:
ERROR_SCHEMA_INSUFFICIENT
</example>

### CRITICAL CONSTRAINTS
1. **NO MARKDOWN**: Do NOT wrap the output in markdown blocks (like ```sql ... ```). Return raw text only.
2. **NO CONVERSATION**: Do NOT include any introductory or concluding text (e.g., "Here is the SQL", "Sure").
3. **NO COMMENTS**: Do NOT add SQL comments (`--`).
4. **STRICT OUTPUT**: The output must be **executable SQL** or the exact error string `ERROR_SCHEMA_INSUFFICIENT`.
"""

logger = logging.getLogger(__name__)


class SqlGenerationOutput(TypedDict):
    query_sql: str
    metadata: MetadataState


def sql_generation_node(state: SqlState, llm: BaseLlm) -> SqlGenerationOutput:
    prompt_template = ChatPromptTemplate.from_messages(
        [
            ("system", system_prompt_sql_generation),
            MessagesPlaceholder("history_messages"),
            ("human", "{user_input}"),
        ]
    )
    chain = prompt_template | llm.langchain
    response = chain.invoke(
        {
            "database_schemas": state.database_schemas,
            "dialect": state.dialect,
            "history_messages": state.history_messages,
            "user_input": state.user_input,
        }
    )
    content = response.content
    text_response = content if isinstance(content, str) else str(content)

    query_sql = SqlClient.extract_sql(text_response)

    logger.info(f"Generated SQL node: {query_sql}")

    return {
        "query_sql": query_sql,
        "metadata": MetadataState(node="sql_generation_node"),
    }
