"""AI 调用工具封装

核心机制：健康门闸（Health Gate）
- ainvoke 抛异常 → 标记故障，全部协程阻塞，单探针退避重试直到成功
- ainvoke 正常返回 → 放行，后续逻辑不变
"""

import asyncio
import re
from datetime import datetime
from typing import Any, Optional

from json_repair import repair_json
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import SystemMessage, HumanMessage
from snowflake import SnowflakeGenerator

from app.core.logger import log

# AI 日志专用雪花ID
_log_id_gen = SnowflakeGenerator(instance=2)

# 匹配 <think>任意内容</think>，用于剥离推理模型的思考过程
_THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL | re.IGNORECASE)

# 匹配 ```json ... ``` 代码块，提取中间的 JSON 内容
_CODE_BLOCK_RE = re.compile(r"```(?:json\w*)?\s*\n?(.*?)\n?\s*```", re.DOTALL | re.IGNORECASE)

# ──────────── 健康门闸 ────────────
_gate_event = asyncio.Event()    # clear=故障阻塞中, set=健康放行
_gate_event.set()                # 初始状态：健康
_gate_probing = False            # 是否已有探针在重试


def parse_llm_json(text: str) -> Any:
    """解析 AI 输出的 JSON，自动去除思考标签、markdown 代码块，容错处理"""
    # 1. 去掉 <think>...</think> 思考内容
    cleaned = _THINK_RE.sub("", text).strip()
    # 2. 如果有 ```json ... ``` 代码块，只取代码块里的内容
    match = _CODE_BLOCK_RE.search(cleaned)
    if match:
        cleaned = match.group(1).strip()
    # 3. repair_json 容错解析：修复不规范的 JSON（多余逗号、缺引号、非法转义等）
    return repair_json(cleaned, return_objects=True)


async def ai_chat(llm: BaseChatModel, system_prompt: str, user_message: str, scene: Optional[str] = None) -> str:
    """异步调用 LLM，返回原始文本。接口异常时触发门闸阻塞重试直到成功。"""
    global _gate_probing

    while True:
        # 门闸检查：如果当前处于故障状态，阻塞等待直到探针恢复门闸
        await _gate_event.wait()

        try:
            messages = [
                SystemMessage(content=system_prompt),
                HumanMessage(content=user_message),
            ]
            response = await llm.ainvoke(messages)
            content = response.content
            # 记录 AI 调用日志（失败不影响主流程）
            if scene:
                await _save_call_log(scene, system_prompt, user_message, content)
            return content
        except Exception as e:
            # 接口失败 → 关闭门闸
            if _gate_event.is_set():
                _gate_event.clear()
                log.error("AI 接口异常，门闸关闭，全部协程阻塞等待恢复: {}", e)

            if not _gate_probing:
                # 当前协程成为探针，负责退避重试直到成功
                _gate_probing = True
                await _probe_until_recover(llm, system_prompt, user_message)
                # 探针恢复后，回到 while 循环顶部用业务数据重新调用
                continue
            # 已有探针在工作，回到 while 顶部 await _gate_event.wait() 等待唤醒


async def _probe_until_recover(llm: BaseChatModel, system_prompt: str, user_message: str) -> str:
    """探针：用轻量 hello 消息探测接口是否恢复，成功后打开门闸，原协程重新走正常流程"""
    global _gate_probing

    delay = 0.2  # 探测间隔 200ms
    attempt = 0

    try:
        while True:
            await asyncio.sleep(delay)
            attempt += 1
            try:
                log.info("探针第{}次探测开始...", attempt)
                messages = [
                    SystemMessage(content="You are a helpful assistant."),
                    HumanMessage(content="hello"),
                ]
                await llm.ainvoke(messages)
                # 成功 → 打开门闸，唤醒所有等待协程
                _gate_probing = False
                _gate_event.set()
                log.info("探针第{}次探测成功，AI 接口恢复，门闸打开", attempt)
                break
            except Exception as e:
                log.warning("探针第{}次探测失败: {}", attempt, e)
    except (asyncio.CancelledError, Exception) as e:
        # 探针意外死亡 → 必须恢复门闸状态，否则所有 worker 永远阻塞
        _gate_probing = False
        _gate_event.set()
        log.error("探针协程意外退出，强制打开门闸: {}", e)

    return ""


async def ai_chat_json(llm: BaseChatModel, system_prompt: str, user_message: str, scene: Optional[str] = None) -> Any:
    """异步调用 LLM，返回解析后的 JSON 对象"""
    raw = await ai_chat(llm, system_prompt, user_message, scene=scene)
    if not raw or not raw.strip():
        log.warning("AI 返回为空")
        return None
    try:
        return parse_llm_json(raw)
    except Exception as e:
        log.warning("AI JSON 解析失败: {}, raw={}", e, raw[:200])
        return None


async def _save_call_log(scene: str, system_prompt: str, user_message: str, response: Optional[str]) -> None:
    """异步写入 AI 调用日志到 PG，失败只 warning 不影响主流程"""
    try:
        from app.core.database import PgSession
        from sqlalchemy import insert
        from app.models.pg.ai_call_log import AiCallLog

        async with PgSession() as pg:
            await pg.execute(
                insert(AiCallLog).values(
                    id=next(_log_id_gen),
                    scene=scene,
                    system_prompt=system_prompt,
                    user_message=user_message,
                    response=response,
                    created_at=datetime.now(),
                )
            )
            await pg.commit()
    except Exception as e:
        log.warning("AI 调用日志写入失败: {}", e)