import logging
from typing import NotRequired, Optional, TypedDict

from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableConfig

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

from ..state import MetadataState, MetadataStatus, SqlState

system_prompt_sql_correction = """
### SYSTEM ROLE
You are a SQL Validation and Correction Engine.
Your task is to fix the provided SQL code based on the specific Error Message given.

### RULES
1. Output ONLY the single, final SQL query as raw text.
2. NO Markdown formatting (do not use ```sql).
3. NO comments, NO explanations, NO conversational filler.
4. Use EXPLICIT JOINS (convert comma-separated tables in FROM clause to INNER/LEFT JOIN).
5. Ensure the query is strictly READ-ONLY (SELECT only).
6. Resolve the ambiguity or syntax error described in the "Error Context".

### INPUT DATA
SQL Code: {sql}
Error Context: {error}

### OUTPUT FORMAT
Return strictly the raw fixed SQL. Do not wrap in Markdown.

### 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**.
"""

logger = logging.getLogger(__name__)


class SqlCorrectionOutput(TypedDict):
    query_sql: Optional[str]
    query_error: Optional[str]
    metadata: MetadataState
    history_messages: NotRequired[list[AnyMessage]]


def sql_correction_node(
    state: SqlState,
    config: RunnableConfig,
    llm: BaseLlm,
) -> SqlCorrectionOutput:
    sql_client: SqlClient = config["configurable"]["sql_client"]
    query_sql = state.query_sql
    query_error = state.query_error

    if not query_sql:
        query_error = "No SQL query to validate"
        return {
            "query_sql": None,
            "query_error": query_error,
            "metadata": MetadataState(
                node="sql_correction_node",
                status=MetadataStatus.ERROR,
                error=query_error,
            ),
        }

    prompt_template = ChatPromptTemplate(("system", system_prompt_sql_correction))
    chain = prompt_template | llm.langchain

    for _ in range(3):
        ok, err = sql_client.validate_sql(query_sql)
        if ok:
            query_sql = query_sql
            query_error = None
            logger.info("Validate SQL node: SQL passed")
            return {
                "query_sql": query_sql,
                "query_error": query_error,
                "metadata": MetadataState(node="sql_correction_node"),
                "history_messages": [
                    HumanMessage(content=state.user_input),
                    AIMessage(content=f"Sql: {query_sql}"),
                ],
            }

        logger.info(f"Validate SQL node: SQL failed: {err}")
        query_error = err

        response = chain.invoke({"sql": query_sql, "error": err}).content
        text_response = response if isinstance(response, str) else str(response)
        query_sql = sql_client.extract_sql(text_response)

    return {
        "query_sql": query_sql,
        "query_error": query_error,
        "metadata": MetadataState(
            node="sql_correction_node",
            status=MetadataStatus.ERROR,
            error=query_error,
        ),
    }
