#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""作业 4 · 成本熔断核心库（只在 207 /opt/workbuddy 运行）。

职责：
1) 单次会话 Token 上限：写入 Dify Agent 快照 max_tokens，并按 conversation 累计拦截。
2) 每日 Token 预算告警：从 Langfuse traces 汇总当日用量，写告警文件。
3) 超阈值模拟降级：不换真实模型供应商（本环境只有 Qwen3.6），改 max_tokens + 提示词标记廉价模型。
"""
from __future__ import annotations

import json
import os
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any
from urllib.parse import quote

import _bootstrap  # noqa: F401
import _agent_cfg as AC
import _config as C
from dify_api import Dify

TZ8 = timezone(timedelta(hours=8))
MARKER = "【成本熔断·预算模式】"
DATA_DIR = Path(getattr(C, "HOMEWORK4_DATA_DIR", "/opt/workbuddy/data/homework4"))


def _ensure_dir() -> Path:
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    return DATA_DIR


def cfg_thresholds() -> dict[str, Any]:
    return {
        "session_max_tokens": int(C.HOMEWORK4_SESSION_MAX_TOKENS),
        "session_token_cap": int(C.HOMEWORK4_SESSION_TOKEN_CAP),
        "daily_token_budget": int(C.HOMEWORK4_DAILY_TOKEN_BUDGET),
        "alert_ratio": float(C.HOMEWORK4_ALERT_RATIO),
        "downgrade_ratio": float(C.HOMEWORK4_DOWNGRADE_RATIO),
        "downgrade_max_tokens": int(C.HOMEWORK4_DOWNGRADE_MAX_TOKENS),
        "normal_model": C.HOMEWORK4_NORMAL_MODEL,
        "budget_model": C.HOMEWORK4_BUDGET_MODEL,
        "app_id": C.HOMEWORK1_APP_ID,
        "agent_id": C.HOMEWORK1_AGENT_ID,
        "snap_id": C.HOMEWORK1_SNAP_ID,
        "draft_id": C.HOMEWORK1_DRAFT_ID,
    }


def save_json(name: str, obj: Any) -> Path:
    p = _ensure_dir() / name
    p.write_text(json.dumps(obj, ensure_ascii=False, indent=2), encoding="utf-8")
    return p


def load_json(name: str, default: Any) -> Any:
    p = DATA_DIR / name
    if not p.is_file():
        return default
    return json.loads(p.read_text(encoding="utf-8"))


def append_alert(kind: str, message: str, extra: dict[str, Any] | None = None) -> Path:
    rec = {
        "ts": datetime.now(TZ8).isoformat(),
        "kind": kind,
        "message": message,
        "extra": extra or {},
    }
    p = _ensure_dir() / "alerts.jsonl"
    with p.open("a", encoding="utf-8") as f:
        f.write(json.dumps(rec, ensure_ascii=False) + "\n")
    save_json("last_alert.json", rec)
    print("[ALERT]", kind, message)
    return p


def lf_public_get(d: Dify, path: str) -> Any:
    url = path if path.startswith("http") else C.LANGFUSE_URL.rstrip("/") + path
    out, err = d.sh(
        "curl -s -m 25 -u '%s:%s' '%s'" % (C.LANGFUSE_PK, C.LANGFUSE_SK, url)
    )
    if err and not out:
        raise RuntimeError("langfuse curl failed: %s" % err)
    try:
        return json.loads(out)
    except json.JSONDecodeError:
        raise RuntimeError("langfuse 非 JSON: %s" % (out[:400],))


def _usage_total(obj: dict[str, Any]) -> int:
    if not isinstance(obj, dict):
        return 0
    usage = obj.get("usage") or obj.get("usageDetails") or obj.get("usage_details") or {}
    if isinstance(usage, dict):
        for k in ("totalTokens", "total_tokens", "total"):
            v = usage.get(k)
            if isinstance(v, (int, float)):
                return int(v)
    for k in ("totalTokens", "total_tokens"):
        v = obj.get(k)
        if isinstance(v, (int, float)):
            return int(v)
    return 0


def list_traces(d: Dify, limit: int = 50) -> list[dict[str, Any]]:
    data = lf_public_get(d, "/api/public/traces?limit=%d" % limit)
    items = data.get("data") if isinstance(data, dict) else data
    return list(items or [])


def observations_tokens(d: Dify, trace_id: str) -> int:
    obs = lf_public_get(d, "/api/public/observations?traceId=%s" % quote(trace_id))
    items = obs.get("data") if isinstance(obs, dict) else []
    total = 0
    for it in items or []:
        if not isinstance(it, dict):
            continue
        typ = (it.get("type") or "").upper()
        if typ in ("GENERATION", "SPAN") or (it.get("name") or "") == "llm":
            total += _usage_total(it)
    return total


def trace_tokens(d: Dify, tr: dict[str, Any]) -> int:
    n = _usage_total(tr)
    if n:
        return n
    tid = tr.get("id") or ""
    if tid:
        return observations_tokens(d, str(tid))
    return 0


def conversation_id_of(tr: dict[str, Any]) -> str:
    md = tr.get("metadata") or {}
    if isinstance(md, dict):
        for k in ("conversation_id", "conversationId"):
            if md.get(k):
                return str(md.get(k))
    inp = tr.get("input")
    text = json.dumps(inp, ensure_ascii=False) if inp is not None else ""
    return ""


def today_range() -> tuple[datetime, datetime]:
    now = datetime.now(TZ8)
    start = now.replace(hour=0, minute=0, second=0, microsecond=0)
    return start, now


def parse_ts(s: str | None) -> datetime | None:
    if not s:
        return None
    try:
        if s.endswith("Z"):
            s = s[:-1] + "+00:00"
        dt = datetime.fromisoformat(s.replace("Z", "+00:00"))
        if dt.tzinfo is None:
            dt = dt.replace(tzinfo=timezone.utc)
        return dt.astimezone(TZ8)
    except Exception:
        return None


def daily_usage(d: Dify) -> dict[str, Any]:
    start, now = today_range()
    used = 0
    traces: list[dict[str, Any]] = []
    for tr in list_traces(d, 80):
        ts = parse_ts(tr.get("timestamp") or tr.get("createdAt") or tr.get("startTime"))
        if ts is None or ts < start:
            continue
        tok = trace_tokens(d, tr)
        used += tok
        traces.append({"id": tr.get("id"), "tokens": tok, "ts": ts.isoformat()})
    return {"date": start.date().isoformat(), "tokens": used, "traces": traces, "until": now.isoformat()}


def session_usage(d: Dify, conversation_id: str) -> int:
    if not conversation_id:
        return 0
    ledger = load_json("session_ledger.json", {})
    ids = set(ledger.get(conversation_id) or [])
    used = 0
    for tr in list_traces(d, 80):
        tid = str(tr.get("id") or "")
        if tid and tid in ids:
            used += trace_tokens(d, tr)
            continue
        cid = conversation_id_of(tr)
        if cid == conversation_id:
            used += trace_tokens(d, tr)
            continue
        blob = json.dumps(tr.get("input") or {}, ensure_ascii=False)
        if conversation_id in blob:
            used += trace_tokens(d, tr)
    return used


def backup_snapshots(d: Dify, snap_id: str, draft_id: str) -> Path:
    snap = AC.load_cfg(d, "agent_config_snapshots", snap_id)
    draft = AC.load_cfg(d, "agent_config_drafts", draft_id)
    return save_json(
        "snapshot_backup.json",
        {"snap_id": snap_id, "draft_id": draft_id, "snap": snap, "draft": draft},
    )


def apply_max_tokens(d: Dify, snap_id: str, draft_id: str, value: int) -> None:
    AC.set_max_tokens(d, snap_id, draft_id, value)


def apply_downgrade(d: Dify, snap_id: str, draft_id: str) -> dict[str, Any]:
    th = cfg_thresholds()
    snap = AC.load_cfg(d, "agent_config_snapshots", snap_id)
    draft = AC.load_cfg(d, "agent_config_drafts", draft_id)

    def patch(cfg: dict) -> dict:
        prompt = cfg.setdefault("prompt", {})
        if not isinstance(prompt, dict):
            prompt = {}
            cfg["prompt"] = prompt
        sys_p = prompt.get("system_prompt") or ""
        line = (
            "%s 已从 %s 降级为 %s（模拟廉价模型，仍走同一 vLLM 端点；"
            "本环境没有第二套商用模型可切）。max_tokens=%d。"
            % (MARKER, th["normal_model"], th["budget_model"], th["downgrade_max_tokens"])
        )
        if MARKER not in sys_p:
            prompt["system_prompt"] = line + "\n\n" + sys_p
        model = cfg.get("model")
        if not isinstance(model, dict):
            model = {}
            cfg["model"] = model
        # 真实调用仍用 Qwen，避免 Dify 找不到模型名导致 400
        model["model"] = th["normal_model"]
        ms = model.get("model_settings")
        if not isinstance(ms, dict):
            ms = {}
            model["model_settings"] = ms
        ms["max_tokens"] = th["downgrade_max_tokens"]
        meta = cfg.setdefault("metadata", {}) if isinstance(cfg.get("metadata"), dict) else {}
        if not isinstance(cfg.get("metadata"), dict):
            cfg["metadata"] = meta
        meta["cost_breaker"] = {
            "mode": "budget",
            "from_model": th["normal_model"],
            "to_model_simulated": th["budget_model"],
            "max_tokens": th["downgrade_max_tokens"],
        }
        return cfg

    AC.save_cfg(d, "agent_config_snapshots", snap_id, patch(snap), "hw4s")
    AC.save_cfg(d, "agent_config_drafts", draft_id, patch(draft), "hw4d")
    rec = {
        "applied": True,
        "from_model": th["normal_model"],
        "to_model_simulated": th["budget_model"],
        "max_tokens": th["downgrade_max_tokens"],
        "ts": datetime.now(TZ8).isoformat(),
    }
    save_json("downgrade.json", rec)
    append_alert("downgrade", "日预算超阈值，已模拟切换到 %s" % th["budget_model"], rec)
    return rec


def restore_snapshots(d: Dify) -> bool:
    bak = load_json("snapshot_backup.json", None)
    if not bak:
        return False
    AC.save_cfg(d, "agent_config_snapshots", bak["snap_id"], bak["snap"], "hw4rs")
    AC.save_cfg(d, "agent_config_drafts", bak["draft_id"], bak["draft"], "hw4rd")
    save_json("downgrade.json", {"applied": False, "restored": True})
    print("[INFO] 已从 snapshot_backup.json 恢复作业 1 Agent 快照")
    return True


def evaluate(d: Dify, conversation_id: str = "") -> dict[str, Any]:
    th = cfg_thresholds()
    daily = daily_usage(d)
    sess = session_usage(d, conversation_id) if conversation_id else 0
    alert_at = int(th["daily_token_budget"] * th["alert_ratio"])
    down_at = int(th["daily_token_budget"] * th["downgrade_ratio"])
    daily_tokens = int(daily["tokens"])
    actions: list[str] = []
    if daily_tokens >= alert_at:
        append_alert(
            "daily_budget",
            "当日 Token=%d 已达日预算 %d 的 %.0f%% 告警线（%d）"
            % (daily_tokens, th["daily_token_budget"], th["alert_ratio"] * 100, alert_at),
            {"daily_tokens": daily_tokens, "budget": th["daily_token_budget"]},
        )
        actions.append("alert")
    if daily_tokens >= down_at:
        already = load_json("downgrade.json", {})
        if already.get("applied"):
            actions.append("downgrade")
        else:
            apply_downgrade(d, th["snap_id"], th["draft_id"])
            actions.append("downgrade")
    session_block = bool(conversation_id) and sess >= th["session_token_cap"]
    if session_block:
        append_alert(
            "session_cap",
            "会话 %s 累计 Token=%d ≥ 上限 %d，拦截后续请求"
            % (conversation_id, sess, th["session_token_cap"]),
            {"conversation_id": conversation_id, "session_tokens": sess},
        )
        actions.append("session_block")
    report = {
        "thresholds": th,
        "daily": daily,
        "session_tokens": sess,
        "conversation_id": conversation_id,
        "alert_at": alert_at,
        "downgrade_at": down_at,
        "session_block": session_block,
        "actions": actions,
    }
    save_json("last_report.json", report)
    return report