"""AI Agent 分析模块 - 调用 OpenAI 兼容接口分析诊断结果 (支持流式输出)"""
import json
import re
import uuid
from collections import Counter
from typing import Optional
import httpx
from backend.config import AI_API_BASE, AI_API_KEY, AI_MODEL
from backend.agent_tools import classify_agent_command, execute_command

SYSTEM_PROMPT = """你是一位资深的 Kubernetes 运维专家和 SRE 工程师。你的任务是分析 K8S 集群诊断报告，找出根本原因并给出修复建议。

分析要求：
1. **问题定位**：明确指出每个问题的根本原因（Root Cause），不要只描述表面现象
2. **影响评估**：说明每个问题对业务的潜在影响
3. **修复方案**：给出具体的修复命令或操作步骤（kubectl 命令、YAML 修改等）
4. **优先级排序**：按紧急程度排序，先处理影响最大的问题
5. **预防建议**：给出避免类似问题再次发生的建议

输出格式：使用清晰的 Markdown，包含标题、列表、代码块。语言使用中文。"""

_HEADERS = {
    "Authorization": f"Bearer {AI_API_KEY}",
    "Content-Type": "application/json",
}


AGENT_SYSTEM_PROMPT = """你是 Kubernetes 自主诊断 Agent。你的首要职责是主动调用命令获取现场证据，而不是只给用户一段脚本或操作建议。

可用工具：
1. kubectl 命令 - 查询集群内资源状态。例如: kubectl get pods -A, kubectl describe node <name>, kubectl logs <pod> -n <ns>
2. 系统诊断 shell 命令 - 查询节点宿主机状态。例如: systemctl status kubelet, journalctl -u kubelet -n 100 --no-pager, crictl ps, df -h, free -m, ss -tlnp, ip addr, top, du -sh, curl, tcpdump
3. containerd 命令 - 查询容器运行时状态。例如: ctr containers list, ctr images list, ctr tasks ls, nerdctl ps -a, nerdctl logs <container>, nerdctl stats
4. etcd 诊断命令 - 检查 etcd 集群健康与数据。例如: etcdctl endpoint status, etcdctl endpoint health, etcdctl member list, etcdctl alarm list, etcdctl get /kubernetes --prefix
5. Linux 基础调优命令 - 例如: iostat -xz 1 3, vmstat 1 3, mpstat -P ALL 1 3, pidstat -d 1 3, lsof -i :<port>, ethtool <iface>, tcpdump -i any -c 10 -nn

行为要求：
1. 如果证据不足，必须返回 command 动作并逐条调查；至少检查与故障资源直接相关的状态、describe、events 或 logs。
2. 禁止在 final 中仅提供"请执行以下命令"的脚本。只有在已经依据真实命令输出形成结论时才能返回 final。
3. 每次只提出一条命令，服务端会执行并把真实输出返回给你，然后你继续下一轮判断。

安全规则：
1. 可以使用管道 (|)、重定向 (>/>>)、逻辑连接符 (&&/||) 或分号 (;) 组合多条命令进行精确排查。不得使用反引号 (`` ` ``) 执行命令替换。
2. kubectl 的 get、describe、logs、top 等只读命令可以自动执行。
3. 系统诊断只读命令 (systemctl status、journalctl、crictl ps、df、free、ss、ip addr 等) 可以自动执行。
4. apply、delete、patch、scale、rollout、exec 等修改或交互命令必须暂停并请求人工批准。
5. systemctl restart/stop/start 等系统修改命令必须暂停并请求人工批准。
6. 不要猜测命令输出；必须依据实际工具结果继续判断。
7. 只要问题没有完全解决（verified=true 且 remaining_issues=0），就要持续执行命令诊断，没有轮数上限。

每轮必须只返回一个 JSON 对象，不要使用 Markdown 代码块：
- 继续调查：{"action":"command","command":"kubectl ...","reason":"为什么执行"}
- 完成分析：{"action":"final","analysis":"中文 Markdown 分析结论","verified":true,"remaining_issues":0}

只有完成修复后的全量复检，并确认 remaining_issues 为 0 时才能返回 final。只定位根因、执行命令成功或部分资源恢复都不能结束。
"""

_PENDING_APPROVALS: dict[str, dict] = {}


def _parse_agent_action(content: str) -> dict:
    text = content.strip()
    candidates = [text]
    fenced = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL | re.IGNORECASE)
    if fenced:
        candidates.insert(0, fenced.group(1))
    first_brace = text.find("{")
    last_brace = text.rfind("}")
    if first_brace >= 0 and last_brace > first_brace:
        candidates.append(text[first_brace:last_brace + 1])

    action = None
    for candidate in candidates:
        try:
            action = json.loads(candidate)
            break
        except json.JSONDecodeError:
            continue
    if action is None:
        raise ValueError("AI 未返回合法 JSON 动作")
    if not isinstance(action, dict) or action.get("action") not in {"command", "final"}:
        raise ValueError("AI 返回了不支持的动作")
    if action["action"] == "command" and not action.get("command"):
        raise ValueError("AI 命令动作缺少 command")
    if action["action"] == "final" and not action.get("analysis"):
        raise ValueError("AI 最终动作缺少 analysis")
    if action["action"] == "final":
        if action.get("verified") is not True or action.get("remaining_issues") != 0:
            raise ValueError("AI 最终动作必须确认 verified=true 且 remaining_issues=0")
    return action


async def _complete_chat(messages: list[dict]) -> str:
    timeout = httpx.Timeout(90, connect=10)
    async with httpx.AsyncClient(timeout=timeout) as client:
        response = await client.post(
            f"{AI_API_BASE}/chat/completions",
            headers=_HEADERS,
            json={
                "model": AI_MODEL,
                "messages": messages,
                "temperature": 0.1,
                "max_tokens": 4096,
                "stream": False,
            },
        )
        response.raise_for_status()
        payload = response.json()
        return payload["choices"][0]["message"]["content"]


async def run_diagnostic_agent(
    report: str,
    question: str = "",
    max_steps: int = 30,
    execution_mode: str = "ai",
    run_id: str = "",
    history_store=None,
):
    """自主诊断循环：AI 模式执行只读命令，手动模式等待用户执行。"""
    if execution_mode not in {"ai", "manual"}:
        yield {"type": "error", "message": "不支持的命令执行模式，仅支持 ai 或 manual"}
        return

    user_content = f"当前 K8S 诊断报告：\n\n{report}"
    if question:
        user_content += f"\n\n用户关注点：{question}"
    messages = [
        {"role": "system", "content": AGENT_SYSTEM_PROMPT},
        {"role": "user", "content": user_content},
    ]

    command_counts = Counter()

    if execution_mode == "ai":
        initial_command = "kubectl get nodes -o wide"
        command_counts[initial_command] += 1
        yield {
            "type": "command", "step": 0, "command": initial_command,
            "reason": "建立集群节点状态基线，确保 AI 执行模式实际采集现场证据",
            "mode": "read",
        }
        initial_result = await execute_command(initial_command)
        yield {"type": "result", **initial_result}
        messages.append({
            "role": "user",
            "content": "AI 执行模式已自动采集第一份现场证据：\n" + json.dumps(initial_result, ensure_ascii=False),
        })

    async for event in _run_agent_loop(
        messages=messages,
        execution_mode=execution_mode,
        max_steps=max_steps,
        start_step=1,
        command_counts=command_counts,
        report=report,
        question=question,
        run_id=run_id,
        history_store=history_store,
    ):
        yield event


async def _run_agent_loop(
    messages: list[dict],
    execution_mode: str,
    max_steps: int,
    start_step: int,
    command_counts: Counter,
    report: str,
    question: str,
    run_id: str = "",
    history_store=None,
):
    invalid_response_count = 0
    step = start_step
    while True:
        if max_steps > 0 and step > max_steps:
            break
        yield {"type": "thinking", "step": step, "message": f"AI 正在规划第 {step} 轮诊断"}
        try:
            content = await _complete_chat(messages)
            action = _parse_agent_action(content)
        except ValueError as exc:
            if "最终动作必须确认" in str(exc):
                invalid_response_count = 0
                messages.extend([
                    {"role": "assistant", "content": content},
                    {
                        "role": "user",
                        "content": (
                            "不能结束诊断：尚未通过修复后的全量复检确认所有问题清零。"
                            "请继续返回 command，检查节点、Pod、Deployment、事件以及原故障资源；"
                            "只有 verified=true 且 remaining_issues=0 时才可返回 final。"
                        ),
                    },
                ])
                yield {"type": "verification_required", "message": "尚未确认所有问题清零，Agent 继续执行诊断验证"}
                continue
            invalid_response_count += 1
            if invalid_response_count >= 3:
                yield {
                    "type": "error",
                    "message": "AI 连续 3 次未返回合法 JSON 动作，已停止本次诊断以避免无效循环",
                }
                return
            messages.extend([
                {"role": "assistant", "content": content},
                {
                    "role": "user",
                    "content": (
                        f"上一次响应无法解析：{exc}。请不要输出解释、脚本或多条命令，"
                        "严格只返回一个 JSON 对象："
                        '{"action":"command","command":"kubectl ...","reason":"..."} '
                        "或在全量复检问题清零后返回 final。"
                    ),
                },
            ])
            yield {
                "type": "response_rejected",
                "message": f"AI 响应格式无效，正在自动纠正并重试（{invalid_response_count}/3）",
                "reason": str(exc),
            }
            continue
        except Exception as exc:
            yield {"type": "error", "message": str(exc)}
            return

        if action["action"] == "final":
            invalid_response_count = 0
            yield {
                "type": "final", "content": action["analysis"], "model": AI_MODEL,
                "verified": True, "remaining_issues": 0,
            }
            return

        command = action["command"]
        invalid_response_count = 0
        reason = action.get("reason", "AI 需要更多集群证据")
        command_counts[command] += 1
        if command_counts[command] > 2:
            messages.extend([
                {"role": "assistant", "content": content},
                {
                    "role": "user",
                    "content": (
                        f"命令 {command!r} 已重复执行两次，不能再次执行。"
                        "请基于已有输出选择不同的只读诊断命令，或在证据充分时返回 final。"
                    ),
                },
            ])
            yield {
                "type": "command_rejected", "command": command,
                "reason": "相同命令已执行两次，请更换诊断方向",
            }
            continue
        try:
            decision = classify_agent_command(command)
        except ValueError as exc:
            messages.extend([
                {"role": "assistant", "content": content},
                {"role": "user", "content": f"命令被安全策略拒绝：{exc}。请改用单条合法 kubectl 或系统诊断命令。"},
            ])
            yield {"type": "command_rejected", "command": command, "reason": str(exc)}
            continue

        yield {
            "type": "command", "step": step, "command": command,
            "reason": reason, "mode": decision.mode,
        }
        if execution_mode == "manual":
            yield {
                "type": "manual_command", "command_id": uuid.uuid4().hex,
                "command": command, "reason": reason, "mode": decision.mode,
                "message": "请在终端手动执行命令，并根据结果继续排查",
            }
            # 手动模式也保存上下文，支持断点续传
            self_save_context(run_id, history_store, messages, command_counts, step, step + 1,
                              max_steps, report, question, execution_mode, content)
            return

        if decision.requires_approval:
            approval_id = uuid.uuid4().hex
            _PENDING_APPROVALS[approval_id] = {
                "command": command,
                "reason": reason,
                "created_by": AI_MODEL,
                "run_id": run_id,
                "report": report,
                "question": question,
                "execution_mode": execution_mode,
                "history_store": history_store,
                "messages": messages + [{"role": "assistant", "content": content}],
                "next_step": step + 1,
                "max_steps": max_steps,
                "command_counts": dict(command_counts),
            }
            # 保存上下文到 DB，避免浏览器关闭/刷新丢失进度
            self_save_context(run_id, history_store, messages, command_counts, step, step + 1,
                              max_steps, report, question, execution_mode, content)
            yield {
                "type": "approval_required", "approval_id": approval_id,
                "command": command, "reason": reason,
            }
            return

        result = await execute_command(command)
        yield {"type": "result", **result}
        messages.extend([
            {"role": "assistant", "content": content},
            {"role": "user", "content": "命令实际执行结果：\n" + json.dumps(result, ensure_ascii=False)},
        ])

        # 每次执行后保存上下文，支持断点续传
        self_save_context(run_id, history_store, messages, command_counts, step, step + 1,
                          max_steps, report, question, execution_mode)

        step += 1

    if max_steps > 0:
        yield {
            "type": "error",
            "message": (
                f"自主诊断已达到安全上限（{max_steps} 轮），但 Agent 仍未形成有效结论。"
                "系统已停止继续执行以避免无限循环，请检查是否重复执行相同命令或 AI 服务响应异常。"
            ),
        }
    # max_steps=0 表示不限轮数，不输出错误 — Agent 会持续执行直到完成或遇到致命错误


async def resume_diagnostic_agent(pending: dict, result: dict):
    """审批命令执行后，把真实结果交回原 Agent 上下文并继续诊断。"""
    messages = list(pending.get("messages") or [])
    messages.append({
        "role": "user",
        "content": (
            "用户已批准命令，以下是实际执行结果。请继续自主诊断并验证问题是否解决：\n"
            + json.dumps(result, ensure_ascii=False)
        ),
    })
    run_id = pending.get("run_id", "")
    history_store = pending.get("history_store")
    async for event in _run_agent_loop(
        messages=messages,
        execution_mode=pending.get("execution_mode", "ai"),
        max_steps=int(pending.get("max_steps", 30)),
        start_step=int(pending.get("next_step", 1)),
        command_counts=Counter(pending.get("command_counts") or {}),
        report=pending.get("report", ""),
        question=pending.get("question", ""),
        run_id=run_id,
        history_store=history_store,
    ):
        yield event


def take_pending_approval(approval_id: str) -> dict | None:
    return _PENDING_APPROVALS.pop(approval_id, None)


def self_save_context(
    run_id: str, history_store,
    messages: list, command_counts: Counter,
    step: int, next_step: int,
    max_steps: int, report: str, question: str,
    execution_mode: str,
    last_response_content: str = "",
):
    """保存 Agent 执行上下文到 DB，支持断点续传。"""
    if not run_id or not history_store:
        return
    context = {
        "messages": messages,
        "command_counts": dict(command_counts),
        "step": step,
        "next_step": next_step,
        "max_steps": max_steps,
        "report": report,
        "question": question,
        "execution_mode": execution_mode,
        "last_response": last_response_content,
    }
    history_store.update_context(run_id, context)


async def resume_from_context(
    ctx: dict,
    run_id: str = "",
    history_store=None,
):
    """从 DB 保存的上下文恢复 Agent 诊断，不重跑历史命令，直接进入下一轮。"""
    messages = list(ctx.get("messages") or [])
    command_counts = Counter(ctx.get("command_counts") or {})
    next_step = int(ctx.get("next_step", 1))
    max_steps = int(ctx.get("max_steps", 0))
    report = ctx.get("report", "")
    question = ctx.get("question", "")
    execution_mode = ctx.get("execution_mode", "ai")

    async for event in _run_agent_loop(
        messages=messages,
        execution_mode=execution_mode,
        max_steps=max_steps,
        start_step=next_step,
        command_counts=command_counts,
        report=report,
        question=question,
        run_id=run_id,
        history_store=history_store,
    ):
        yield event


async def _stream_chat(messages: list[dict]):
    """通用流式请求 - yield 文本 chunk"""
    # connect=10s 快速失败; read=90s 防止 AI 服务无响应时前端无限卡死
    timeout = httpx.Timeout(90, connect=10)
    async with httpx.AsyncClient(timeout=timeout) as client:
        async with client.stream(
            "POST",
            f"{AI_API_BASE}/chat/completions",
            headers=_HEADERS,
            json={
                "model": AI_MODEL,
                "messages": messages,
                "temperature": 0.3,
                "max_tokens": 4096,
                "stream": True,
            },
        ) as resp:
            # 非 2xx 直接抛异常 (由上层转成 error 事件)
            resp.raise_for_status()

            # 有些服务认证失败时返回 200 + JSON 错误体 (非 SSE)，需要识别
            ctype = resp.headers.get("content-type", "")
            if "text/event-stream" not in ctype and "application/json" in ctype:
                body = await resp.aread()
                try:
                    err = json.loads(body.decode())
                    raise httpx.HTTPStatusError(
                        f"AI 服务返回错误: {err.get('error', body.decode()[:300])}",
                        request=resp.request,
                        response=resp,
                    )
                except (json.JSONDecodeError, UnicodeDecodeError):
                    raise httpx.HTTPStatusError(
                        f"AI 服务返回非流式响应: {body.decode()[:300]}",
                        request=resp.request,
                        response=resp,
                    )

            async for line in resp.aiter_lines():
                if not line.startswith("data: "):
                    continue
                data_str = line[6:].strip()
                if data_str == "[DONE]":
                    break
                try:
                    chunk = json.loads(data_str)
                    delta = chunk["choices"][0].get("delta", {})
                    content = delta.get("content", "")
                    if content:
                        yield content
                except (json.JSONDecodeError, KeyError, IndexError):
                    continue


async def analyze_with_ai_stream(report: str, question: str = ""):
    """流式 AI 分析诊断报告 - yield 文本 chunk"""
    user_msg = f"以下是 K8S 集群诊断报告，请进行详细分析：\n\n{report}"
    if question:
        user_msg += f"\n\n用户额外问题：{question}"

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_msg},
    ]
    async for chunk in _stream_chat(messages):
        yield chunk


async def chat_with_ai_stream(messages: list[dict], context: str = ""):
    """流式多轮对话 - yield 文本 chunk"""
    msgs = [{"role": "system", "content": SYSTEM_PROMPT}]
    if context:
        msgs.append({"role": "system", "content": f"当前诊断上下文：\n{context}"})
    msgs.extend(messages)
    async for chunk in _stream_chat(msgs):
        yield chunk


async def analyze_with_ai(report: str, question: str = "") -> dict:
    """非流式 AI 分析 (兼容旧接口)"""
    try:
        content = ""
        async for chunk in analyze_with_ai_stream(report, question):
            content += chunk
        return {"success": True, "analysis": content, "model": AI_MODEL}
    except httpx.HTTPStatusError as e:
        return {"success": False, "analysis": f"AI 接口返回错误: {e.response.status_code} - {e.response.text[:500]}", "model": AI_MODEL}
    except httpx.ConnectError:
        return {"success": False, "analysis": f"无法连接 AI 服务 ({AI_API_BASE})，请检查 .env 中的 AI_API_BASE 配置", "model": AI_MODEL}
    except Exception as e:
        return {"success": False, "analysis": f"AI 分析失败: {str(e)}", "model": AI_MODEL}


async def chat_with_ai(messages: list[dict], context: str = "") -> dict:
    """非流式多轮对话 (兼容旧接口)"""
    try:
        content = ""
        async for chunk in chat_with_ai_stream(messages, context):
            content += chunk
        return {"success": True, "reply": content}
    except Exception as e:
        return {"success": False, "reply": f"AI 对话失败: {str(e)}"}