"""NL2SQL SQL 生成器：只负责调用 Chat 模型并清理模型输出。"""
from __future__ import annotations

import json
import re
from dataclasses import dataclass
from typing import Any

from sqlglot import exp, parse
from sqlglot.errors import ParseError

from tool.llm import llm
from nl2sql.semantics import build_semantic_context


class SqlGenerationError(ValueError):
    """模型未返回可用的只读 SQL。"""


@dataclass(frozen=True)
class GeneratedSql:
    sql: str


def _clean_model_output(output: str) -> str:
    value = (output or "").strip()
    value = re.sub(r"^```(?:sql)?\s*", "", value, flags=re.IGNORECASE)
    value = re.sub(r"\s*```$", "", value).strip()
    return value.removesuffix(";").strip()


def _build_prompt(
    question: str,
    schema: dict[str, Any],
    few_shot: list[dict[str, Any]],
    semantic_context: dict[str, Any] | None = None,
    conversation_context: str = "",
) -> str:
    return (
        "你是基金业务数据库 SQL 生成器。\n"
        "只根据提供的 Schema 生成一条 MySQL SELECT，禁止写操作、跨库访问和未提供的表字段。\n"
        "只输出 SQL，不要输出解释、Markdown 或代码围栏。\n"
        f"用户问题：{question}\n"
        f"Schema：{json.dumps(schema, ensure_ascii=False, sort_keys=True)}\n"
        f"业务语义：{json.dumps(semantic_context or {}, ensure_ascii=False, sort_keys=True)}\n"
        f"会话上下文：{conversation_context}\n"
        f"Few-shot：{json.dumps(few_shot or [], ensure_ascii=False, sort_keys=True)}"
    )


async def generate_sql(
    question: str,
    schema: dict[str, Any],
    *,
    llm_client=llm,
    few_shot: list[dict[str, Any]] | None = None,
    semantic_context: dict[str, Any] | None = None,
    conversation_context: str = "",
) -> GeneratedSql:
    """调用 Chat 模型生成 SQL，并在返回前确认其为单条 SELECT。"""
    if not question or not question.strip():
        raise SqlGenerationError("用户问题不能为空")
    messages = [
        {"role": "system", "content": "你必须严格遵守只输出单条 SELECT SQL。"},
        {
            "role": "user",
            "content": _build_prompt(
                question,
                schema,
                few_shot or [],
                semantic_context or build_semantic_context(question, schema),
                conversation_context,
            ),
        },
    ]
    try:
        output = await llm_client.chat(messages, temperature=0, max_tokens=256)
    except Exception as exc:  # noqa: BLE001 统一收敛模型异常
        raise SqlGenerationError("SQL 模型调用失败") from exc
    sql = _clean_model_output(output)
    if not sql:
        raise SqlGenerationError("模型未返回 SQL")
    try:
        statements = parse(sql, read="mysql")
    except ParseError as exc:
        raise SqlGenerationError("模型返回的 SQL 无法解析") from exc
    if len(statements) != 1 or not isinstance(statements[0], exp.Select):
        raise SqlGenerationError("模型只允许返回单条 SELECT")
    return GeneratedSql(sql=statements[0].sql(dialect="mysql"))