"""
受控 ReAct Text2SQL 工具
核心思想：Agent负责推理，Tool负责SQL，Reviewer负责纠错

内部流程：
1. Schema Injection - 注入数据库表结构
2. SQL Generation - LLM生成SQL（使用 SQLAlchemy 连接池 + extra_info 辅助）
3. SQL Review - LLM审查SQL正确性
4. SQL Guard - 代码级安全校验（硬限制）
5. SQL Execute - 执行SQL并返回结果（使用 SQLAlchemy 连接池单例）
6. Merge - 相似结果合并（Embedding）
"""
import re
import json
import asyncio
from datetime import datetime
from typing import Dict, Any, Optional, List, Tuple

from sqlalchemy import text
from langchain_community.utilities import SQLDatabase

from config import POSTGRES_CONNECTION_STRING,POSTGRES_SQLALCHEMY_URL
from modelsAPI.model_api import OpenaiAPI
from prompts import TEXT2SQL_PROMPTS
from tools.texttosql_utils import get_extra_info

# ============================================================
# 数据库连接池单例（避免重复初始化）
# ============================================================

_db_instance: Optional[SQLDatabase] = None
_db_lock = asyncio.Lock()


async def _get_db() -> SQLDatabase:
    """懒加载并缓存数据库连接，避免重复初始化"""
    global _db_instance
    if _db_instance is None:
        async with _db_lock:
            if _db_instance is None:
                _db_instance = await asyncio.to_thread(
                    SQLDatabase.from_uri, POSTGRES_SQLALCHEMY_URL
                )
    return _db_instance


# ============================================================
# Step 1: Schema Injection
# ============================================================

FAULT_RECORDS_SCHEMA = {
    "table_name": "fault_records",
    "description": "故障维修记录表，存储所有故障诊断和维修的记录",
    "columns": {
        "id": {
            "type": "SERIAL",
            "description": "自增主键",
            "usable_in_query": False
        },
        "ship_number": {
            "type": "TEXT",
            "description": "舰船舷号，如'101'、'163'、'554'等，2-4位数字",
            "usable_in_query": True,
            "examples": ["101", "163", "554", "602"]
        },
        "device_name": {
            "type": "TEXT",
            "description": "设备名称，如'主机'、'发电机'、'分油机'、'空压机'等",
            "usable_in_query": True,
            "examples": ["主机", "发电机", "分油机", "空压机", "舵机", "锅炉"]
        },
        "fault": {
            "type": "TEXT",
            "description": "故障现象描述，如'滑油压力低'、'异常振动'、'启动失败'等",
            "usable_in_query": True,
            "examples": ["滑油压力低", "异常振动", "启动失败", "排烟温度高"]
        },
        "spare_parts": {
            "type": "JSONB",
            "description": "消耗的备件列表，JSON数组格式，如'[\"密封圈\", \"滤芯\"]'",
            "usable_in_query": True,
            "json_query_hint": "使用 jsonb_array_elements_text(spare_parts) 展开备件进行统计"
        },
        "system_name": {
            "type": "TEXT",
            "description": "所属系统名称，如'动力系统'、'电气系统'、'导航系统'等",
            "usable_in_query": True,
            "examples": ["动力系统", "电气系统", "导航系统", "通信系统", "其他"]
        },
        "created_at": {
            "type": "TIMESTAMP",
            "description": "记录创建时间，默认为当前时间",
            "usable_in_query": True,
            "examples": ["2025-01-15 10:30:00"]
        }
    }
}

# 字段白名单（SQL Guard 使用）
ALLOWED_COLUMNS = set(FAULT_RECORDS_SCHEMA["columns"].keys())


def get_schema_prompt_section() -> str:
    """生成 Schema 注入的提示词部分"""
    lines = [f"数据库表：{FAULT_RECORDS_SCHEMA['table_name']}"]
    lines.append(f"说明：{FAULT_RECORDS_SCHEMA['description']}")
    lines.append("")
    lines.append("字段列表：")

    for col_name, col_info in FAULT_RECORDS_SCHEMA["columns"].items():
        if not col_info.get("usable_in_query", True):
            continue
        line = f"  - {col_name} ({col_info['type']}): {col_info['description']}"
        if "examples" in col_info:
            line += f"（示例：{', '.join(col_info['examples'][:3])}）"
        if "json_query_hint" in col_info:
            line += f"\n    查询提示：{col_info['json_query_hint']}"
        lines.append(line)

    return "\n".join(lines)


# ============================================================
# Step 2: SQL Generation
# ============================================================

async def generate_sql(question: str, current_date: str) -> Dict[str, Any]:
    """
    使用 LLM 根据用户问题生成 SQL。
    并发预热数据库连接与获取 extra_info，提升整体响应速度。

    Returns:
        {"sql": "SELECT ...", "success": True/False, "error": "..."}
    """
    

    schema_section = get_schema_prompt_section()

    # 并发：获取辅助信息 + 预热DB连接
    extra_info, _ = await asyncio.gather(
        get_extra_info(question),
        _get_db(),
    )

    prompt = TEXT2SQL_PROMPTS["sql_generation"].format(
        schema=schema_section,
        question=question,
        current_date=current_date,
        extra_info=extra_info or ""
    )

    try:
        result = await OpenaiAPI.open_api_chat_without_thinking(
            prompt,
            model=None,
            json_output=True
        )

        # 提取 JSON
        json_match = re.search(r'\{.*\}', result, re.DOTALL)
        if json_match:
            parsed = json.loads(json_match.group(0))
            sql = parsed.get("sql", "").strip()
            if sql:
                return {"sql": sql, "success": True, "error": ""}

        # 兜底：直接提取 SELECT 语句
        sql_match = re.search(r'(SELECT\s+.+?;?)$', result.strip(), re.IGNORECASE | re.DOTALL)
        if sql_match:
            sql = sql_match.group(1).strip().rstrip(';')
            return {"sql": sql, "success": True, "error": ""}

        return {"sql": "", "success": False, "error": "LLM未返回有效SQL"}

    except Exception as e:
        return {"sql": "", "success": False, "error": f"SQL生成失败: {str(e)}"}


# ============================================================
# Step 3: SQL Review (LLM 审查)
# ============================================================

async def review_sql(question: str, sql: str) -> Dict[str, Any]:
    """
    使用 LLM 审查 SQL 是否真正回答了用户问题

    Returns:
        {
            "approved": True/False,
            "issues": ["问题1", "问题2"],
            "suggested_sql": "修正后的SQL（如果有问题）",
            "review_comment": "审查意见"
        }
    """
    from modelsAPI.model_api import OpenaiAPI
    from prompts import TEXT2SQL_PROMPTS

    schema_section = get_schema_prompt_section()
    prompt = TEXT2SQL_PROMPTS["sql_review"].format(
        schema=schema_section,
        question=question,
        sql=sql
    )

    try:
        result = await OpenaiAPI.open_api_chat_without_thinking(
            prompt,
            model=None,
            json_output=True
        )

        json_match = re.search(r'\{.*\}', result, re.DOTALL)
        if json_match:
            parsed = json.loads(json_match.group(0))
            approved = parsed.get("approved", False)
            issues = parsed.get("issues", [])
            suggested_sql = parsed.get("suggested_sql", "").strip()
            review_comment = parsed.get("review_comment", "")

            if suggested_sql:
                suggested_sql = suggested_sql.rstrip(';')

            return {
                "approved": approved,
                "issues": issues,
                "suggested_sql": suggested_sql,
                "review_comment": review_comment
            }

        return {
            "approved": True,
            "issues": [],
            "suggested_sql": "",
            "review_comment": "审查结果解析失败，默认通过"
        }

    except Exception as e:
        return {
            "approved": True,
            "issues": [],
            "suggested_sql": "",
            "review_comment": f"审查过程异常，默认通过: {str(e)}"
        }


# ============================================================
# Step 4: SQL Guard (代码级安全校验)
# ============================================================

class SQLGuardError(Exception):
    """SQL Guard 校验失败异常"""
    pass


def guard_sql(sql: str) -> Tuple[bool, str, str]:
    """
    SQL 安全守卫 - 代码硬限制

    Returns:
        (passed, cleaned_sql, error_message)
    """
    if not sql or not sql.strip():
        return False, "", "SQL为空"

    sql_upper = sql.strip().upper()

    # 1. 只允许 SELECT
    if not sql_upper.startswith("SELECT"):
        return False, "", "只允许SELECT查询，禁止INSERT/UPDATE/DELETE/CREATE等操作"

    # 2. 禁止 UNION
    if "UNION" in sql_upper:
        return False, "", "禁止使用UNION操作"

    # 3. 禁止子查询嵌套过深（最多2层）
    subquery_count = sql_upper.count("SELECT") - 1
    if subquery_count > 2:
        return False, "", f"子查询嵌套过深（{subquery_count}层），最多允许2层"

    # 4. 禁止危险函数
    dangerous_functions = ["PG_SLEEP", "SLEEP", "BENCHMARK", "WAITFOR", "DELAY"]
    for func in dangerous_functions:
        if func in sql_upper:
            return False, "", f"禁止使用危险函数: {func}"

    # 5. 禁止系统表访问
    system_tables = ["PG_", "INFORMATION_SCHEMA", "PG_CATALOG", "PG_CLASS"]
    for sys_table in system_tables:
        if sys_table in sql_upper:
            return False, "", f"禁止访问系统表: {sys_table}"

    # 6. 禁止注释注入
    if "--" in sql or "/*" in sql or "*/" in sql:
        return False, "", "SQL中不允许包含注释"

    # 7. 禁止分号（防止多语句注入）
    if ";" in sql.rstrip(';'):
        return False, "", "禁止多语句执行"

    # 8. 检查字段白名单
    used_columns = set()
    dot_matches = re.findall(r'fault_records\.(\w+)', sql, re.IGNORECASE)
    used_columns.update(m.lower() for m in dot_matches)

    select_match = re.search(r'SELECT\s+(.*?)\s+FROM', sql, re.IGNORECASE | re.DOTALL)
    if select_match:
        select_part = select_match.group(1)
        if '*' not in select_part:
            for part in select_part.split(','):
                part = part.strip()
                col_match = re.search(r'(\w+)\s*(?:AS|as)?\s*$', part)
                if col_match:
                    col_name = col_match.group(1).lower()
                    if col_name not in ('count', 'sum', 'avg', 'max', 'min', 'as',
                                        'desc', 'asc', 'null', 'distinct', 'all',
                                        'jsonb_array_elements_text', 'jsonb_array_elements',
                                        'coalesce', 'cast', 'date', 'timestamp', 'text'):
                        used_columns.add(col_name)

    for col in used_columns:
        if col not in ALLOWED_COLUMNS and col != 'id':
            return False, "", f"使用了不允许的字段: {col}"

    # 9. 自动添加 LIMIT（如果没有）
    cleaned_sql = sql.strip().rstrip(';')
    if "LIMIT" not in sql_upper:
        cleaned_sql = f"{cleaned_sql} LIMIT 100"

    # 10. 强制只查 fault_records 表
    from_match = re.search(r'FROM\s+(\w+)', sql, re.IGNORECASE)
    if from_match:
        table_name = from_match.group(1).lower()
        if table_name != "fault_records":
            return False, "", f"只允许查询 fault_records 表，不允许查询: {table_name}"

    return True, cleaned_sql, ""


# ============================================================
# Step 5: SQL Execute
# ============================================================

async def execute_sql(sql: str) -> Dict[str, Any]:
    """
    执行 SQL 并返回结果。
    使用模块级 SQLAlchemy 连接池单例，避免每次重复建立连接。

    Returns:
        {"rows": [...], "columns": [...], "row_count": int, "success": True/False, "error": ""}
    """
    try:
        db = await _get_db()

        def _run() -> Dict[str, Any]:
            try:
                with db._engine.connect() as conn:
                    result = conn.execute(text(sql))
                    columns = list(result.keys())
                    result_rows = []
                    for row in result.fetchall():
                        row_dict = {}
                        for i, col in enumerate(columns):
                            val = row[i]
                            if isinstance(val, datetime):
                                val = val.isoformat()
                            elif isinstance(val, (list, dict)):
                                val = json.dumps(val, ensure_ascii=False)
                            elif not isinstance(val, (str, int, float, bool, type(None))):
                                val = str(val)
                            row_dict[col] = val
                        result_rows.append(row_dict)
                    conn.commit()
                    return {
                        "rows": result_rows,
                        "columns": columns,
                        "row_count": len(result_rows),
                        "success": True,
                        "error": ""
                    }
            except Exception as e:
                return {
                    "rows": [],
                    "columns": [],
                    "row_count": 0,
                    "success": False,
                    "error": f"SQL执行失败: {str(e)}"
                }

        return await asyncio.to_thread(_run)

    except Exception as e:
        return {
            "rows": [],
            "columns": [],
            "row_count": 0,
            "success": False,
            "error": f"获取数据库连接失败: {str(e)}"
        }


# ============================================================
# Text2SQL 主入口 - 内部闭环
# ============================================================

MAX_RETRY = 2  # 最大重试次数（生成→审查→修正）


async def text2sql_tool(question: str, current_date: Optional[str] = None) -> Dict[str, Any]:
    """
    受控 Text2SQL 工具 - 内部形成闭环

    流程：question → SQL生成 → SQL审查 → SQL Guard → SQL执行 → 返回结果

    Args:
        question: 用户问题
        current_date: 当前日期（格式：YYYY-MM-DD），默认为今天

    Returns:
        {
            "success": True/False,
            "sql": "最终执行的SQL",
            "rows": [...],
            "columns": [...],
            "row_count": int,
            "error": "",
            "review_log": [...]  # 审查日志
        }
    """
    if not current_date:
        current_date = datetime.now().strftime("%Y-%m-%d")

    review_log = []
    current_sql = ""

    for attempt in range(MAX_RETRY + 1):
        # Step 2: SQL Generation
        if attempt == 0:
            gen_result = await generate_sql(question, current_date)
        else:
            gen_result = await generate_sql_with_feedback(
                question, current_date, current_sql, review_log
            )

        if not gen_result["success"]:
            return {
                "success": False,
                "sql": "",
                "rows": [],
                "columns": [],
                "row_count": 0,
                "error": gen_result["error"],
                "review_log": review_log
            }

        current_sql = gen_result["sql"]

        # Step 3: SQL Review
        review_result = await review_sql(question, current_sql)
        review_log.append({
            "attempt": attempt + 1,
            "sql": current_sql,
            "approved": review_result["approved"],
            "issues": review_result["issues"],
            "comment": review_result["review_comment"]
        })

        if not review_result["approved"]:
            if review_result["suggested_sql"]:
                current_sql = review_result["suggested_sql"]
                review_log.append({
                    "attempt": attempt + 1,
                    "sql": current_sql,
                    "note": "使用审查建议的修正SQL"
                })
            elif attempt < MAX_RETRY:
                continue
            else:
                return {
                    "success": False,
                    "sql": current_sql,
                    "rows": [],
                    "columns": [],
                    "row_count": 0,
                    "error": f"SQL审查不通过: {'; '.join(review_result['issues'])}",
                    "review_log": review_log
                }

        # Step 4: SQL Guard
        passed, cleaned_sql, guard_error = guard_sql(current_sql)
        if not passed:
            review_log.append({
                "attempt": attempt + 1,
                "sql": current_sql,
                "guard_passed": False,
                "guard_error": guard_error
            })
            if attempt < MAX_RETRY:
                review_log.append({"note": f"Guard拒绝: {guard_error}，将重新生成"})
                continue
            else:
                return {
                    "success": False,
                    "sql": current_sql,
                    "rows": [],
                    "columns": [],
                    "row_count": 0,
                    "error": f"SQL安全校验失败: {guard_error}",
                    "review_log": review_log
                }

        current_sql = cleaned_sql
        review_log.append({
            "attempt": attempt + 1,
            "sql": current_sql,
            "guard_passed": True,
            "note": "Guard通过"
        })

        # Step 5: SQL Execute
        exec_result = await execute_sql(current_sql)

        if not exec_result["success"]:
            if attempt < MAX_RETRY:
                review_log.append({
                    "attempt": attempt + 1,
                    "sql": current_sql,
                    "exec_error": exec_result["error"],
                    "note": "执行失败，将重新生成"
                })
                continue
            else:
                return {
                    "success": False,
                    "sql": current_sql,
                    "rows": [],
                    "columns": [],
                    "row_count": 0,
                    "error": exec_result["error"],
                    "review_log": review_log
                }

        # Step 6: 相似结果合并
        merged_rows = await merge_similar_rows(
            exec_result["rows"], exec_result["columns"], question
        )

        return {
            "success": True,
            "sql": current_sql,
            "rows": merged_rows,
            "columns": exec_result["columns"],
            "row_count": len(merged_rows),
            "raw_row_count": exec_result["row_count"],
            "error": "",
            "review_log": review_log
        }

    return {
        "success": False,
        "sql": current_sql,
        "rows": [],
        "columns": [],
        "row_count": 0,
        "error": "超过最大重试次数",
        "review_log": review_log
    }


async def generate_sql_with_feedback(
    question: str,
    current_date: str,
    previous_sql: str,
    review_log: List[Dict[str, Any]]
) -> Dict[str, Any]:
    """
    基于审查反馈重新生成 SQL
    """
    from modelsAPI.model_api import OpenaiAPI
    from prompts import TEXT2SQL_PROMPTS

    schema_section = get_schema_prompt_section()

    feedback_parts = []
    for log_entry in review_log:
        if "issues" in log_entry and log_entry["issues"]:
            feedback_parts.append(f"问题: {'; '.join(log_entry['issues'])}")
        if "guard_error" in log_entry:
            feedback_parts.append(f"安全校验失败: {log_entry['guard_error']}")
        if "exec_error" in log_entry:
            feedback_parts.append(f"执行失败: {log_entry['exec_error']}")

    feedback = "\n".join(feedback_parts) if feedback_parts else "无具体反馈"

    prompt = TEXT2SQL_PROMPTS["sql_regeneration"].format(
        schema=schema_section,
        question=question,
        current_date=current_date,
        previous_sql=previous_sql,
        feedback=feedback
    )

    try:
        result = await OpenaiAPI.open_api_chat_without_thinking(
            prompt,
            model=None,
            json_output=True
        )

        json_match = re.search(r'\{.*\}', result, re.DOTALL)
        if json_match:
            parsed = json.loads(json_match.group(0))
            sql = parsed.get("sql", "").strip()
            if sql:
                return {"sql": sql, "success": True, "error": ""}

        sql_match = re.search(r'(SELECT\s+.+?;?)$', result.strip(), re.IGNORECASE | re.DOTALL)
        if sql_match:
            sql = sql_match.group(1).strip().rstrip(';')
            return {"sql": sql, "success": True, "error": ""}

        return {"sql": "", "success": False, "error": "重新生成SQL失败"}

    except Exception as e:
        return {"sql": "", "success": False, "error": f"重新生成SQL异常: {str(e)}"}


# ============================================================
# Step 6: 相似结果合并（Embedding）
# ============================================================

MERGE_KEYWORDS = ["故障", "频次", "排名", "高频", "最多", "top", "发生次数"]


def _should_merge(columns: List[str], question: str) -> bool:
    """判断是否需要合并相似行"""
    question_lower = question.lower()
    if any(kw in question_lower for kw in MERGE_KEYWORDS):
        col_names = [c.lower() for c in columns]
        has_name = any(n in col_names for n in ["name", "故障", "fault", "device_name"])
        has_count = any(n in col_names for n in ["count", "次数", "频次"])
        return has_name and has_count
    return False


def _find_name_column(columns: List[str]) -> Optional[str]:
    """找到名称列"""
    priority = ["name", "故障名称", "故障", "fault", "device_name", "备件", "spare"]
    col_lower_map = {c.lower(): c for c in columns}
    for p in priority:
        if p in col_lower_map:
            return col_lower_map[p]
    for c in columns:
        if c.lower() not in ("count", "次数", "频次", "rank", "排名", "system_name", "系统"):
            return c
    return None


def _find_count_column(columns: List[str]) -> Optional[str]:
    """找到计数列"""
    priority = ["count", "次数", "频次", "cnt"]
    col_lower_map = {c.lower(): c for c in columns}
    for p in priority:
        if p in col_lower_map:
            return col_lower_map[p]
    return None


def _find_system_column(columns: List[str]) -> Optional[str]:
    """找到系统分类列"""
    priority = ["system_name", "系统", "系统分类", "system"]
    col_lower_map = {c.lower(): c for c in columns}
    for p in priority:
        if p in col_lower_map:
            return col_lower_map[p]
    return None


async def merge_similar_rows(
    rows: List[Dict[str, Any]],
    columns: List[str],
    question: str,
    similarity_threshold: float = 0.85
) -> List[Dict[str, Any]]:
    """
    使用 embedding 合并相似的故障/备件行

    例如：
    - "连杆大端轴承发生损坏异常" (7次, 其他) + "连杆大端轴承发生损坏异常" (6次, 汽车配件)
      → "连杆大端轴承发生损坏异常" (13次, 其他)
    - "发动机中空冷器管束芯体结垢或堵塞..." + "空冷器管束芯体结垢或堵塞..."
      → 合并为一条
    """
    if not rows or len(rows) <= 1:
        return rows

    if not _should_merge(columns, question):
        return rows

    name_col = _find_name_column(columns)
    count_col = _find_count_column(columns)
    system_col = _find_system_column(columns)

    if not name_col or not count_col:
        return rows

    try:
        from modelsAPI.model_api import OpenaiAPI

        names = [row[name_col] for row in rows if row.get(name_col)]
        if len(names) <= 1:
            return rows

        embeddings = await OpenaiAPI.get_embeddings_batch_async(names)
        name_emb_map = {name: emb for name, emb in zip(names, embeddings)}

        sorted_rows = sorted(rows, key=lambda r: int(r.get(count_col, 0)), reverse=True)

        merged_rows = []
        merged_indices = set()

        for i, row in enumerate(sorted_rows):
            if i in merged_indices:
                continue

            name_i = row.get(name_col, "")
            emb_i = name_emb_map.get(name_i)

            if not emb_i:
                merged_rows.append(row)
                continue

            merged_row = dict(row)
            total_count = int(row.get(count_col, 0))
            systems = set()
            if system_col and row.get(system_col):
                systems.add(row[system_col])

            for j in range(i + 1, len(sorted_rows)):
                if j in merged_indices:
                    continue

                name_j = sorted_rows[j].get(name_col, "")
                emb_j = name_emb_map.get(name_j)

                if not emb_j:
                    continue

                similarity = OpenaiAPI.cosine_similarity(emb_i, emb_j)

                if similarity >= similarity_threshold:
                    total_count += int(sorted_rows[j].get(count_col, 0))
                    if system_col and sorted_rows[j].get(system_col):
                        systems.add(sorted_rows[j][system_col])
                    merged_indices.add(j)

            merged_row[count_col] = total_count
            if system_col and systems:
                merged_row[system_col] = "、".join(sorted(systems))

            merged_rows.append(merged_row)
            merged_indices.add(i)

        merged_rows.sort(key=lambda r: int(r.get(count_col, 0)), reverse=True)

        if "rank" in columns or "排名" in columns:
            rank_col = "rank" if "rank" in columns else "排名"
            for idx, row in enumerate(merged_rows):
                row[rank_col] = idx + 1

        return merged_rows

    except Exception as e:
        print(f"Embedding合并失败，使用原始数据: {str(e)}")
        return rows