#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""作业 6 · 自动周报：Langfuse Token + Prometheus GPU + Qwen 总结 + 钉钉推送。"""
from __future__ import annotations

import json
import re
import urllib.parse
import urllib.request
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Callable, Optional

import _bootstrap  # noqa: F401
import _config as C

TZ8 = timezone(timedelta(hours=8))
HTTP_GET = Callable[[str], str]
PROM_QUERIES = (
    "avg_over_time(DCGM_FI_DEV_GPU_UTIL[7d])",
    "avg_over_time(nvidia_gpu_duty_cycle[7d])",
    "avg_over_time(nvidia_smi_utilization_gpu_ratio[7d])",
)


def data_dir() -> Path:
    p = Path(getattr(C, "HOMEWORK6_DATA_DIR", "/opt/workbuddy/data/homework6"))
    p.mkdir(parents=True, exist_ok=True)
    return p


def curl_get(url: str, auth: tuple[str, str] | None = None, timeout: int = 30) -> str:
    req = urllib.request.Request(url)
    if auth:
        import base64

        token = base64.b64encode(("%s:%s" % auth).encode()).decode()
        req.add_header("Authorization", "Basic %s" % token)
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return resp.read().decode("utf-8", "replace")


def curl_post_json(url: str, body: dict[str, Any], timeout: int = 30) -> dict[str, Any]:
    data = json.dumps(body, ensure_ascii=False).encode("utf-8")
    req = urllib.request.Request(
        url, data=data, headers={"Content-Type": "application/json"}, method="POST"
    )
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        raw = resp.read().decode("utf-8", "replace")
    return json.loads(raw)


def gpu_exec(cmd: str, runner: Optional[HTTP_GET] = None, timeout: int = 40) -> str:
    if runner is not None:
        return runner(cmd)
    import paramiko

    c = paramiko.SSHClient()
    c.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    c.connect(C.GPU_HOST, username=C.GPU_USER, password=C.GPU_PASS, timeout=20)
    try:
        _i, o, e = c.exec_command(cmd, timeout=timeout)
        out = o.read().decode("utf-8", "replace")
        err = e.read().decode("utf-8", "replace")
        if not out.strip() and err.strip():
            raise RuntimeError(err.strip())
        return out
    finally:
        c.close()


def parse_langfuse_daily(raw: str, days: int = 7) -> dict[str, Any]:
    data = json.loads(raw)
    rows = data.get("data") if isinstance(data, dict) else []
    daily: list[dict[str, Any]] = []
    total_tokens = 0
    total_cost = 0.0
    total_traces = 0
    for row in rows or []:
        if not isinstance(row, dict):
            continue
        day_tokens = 0
        models: list[dict[str, Any]] = []
        for u in row.get("usage") or []:
            if not isinstance(u, dict):
                continue
            tok = int(u.get("totalUsage") or 0)
            day_tokens += tok
            if u.get("model"):
                models.append(
                    {
                        "model": u.get("model"),
                        "tokens": tok,
                        "cost": float(u.get("totalCost") or 0),
                    }
                )
        total_tokens += day_tokens
        total_cost += float(row.get("totalCost") or 0)
        total_traces += int(row.get("countTraces") or 0)
        daily.append(
            {
                "date": row.get("date"),
                "tokens": day_tokens,
                "traces": int(row.get("countTraces") or 0),
                "cost": float(row.get("totalCost") or 0),
                "models": models,
            }
        )
    daily.sort(key=lambda x: x.get("date") or "")
    return {
        "days": days,
        "daily": daily[-days:],
        "total_tokens": total_tokens,
        "total_cost": round(total_cost, 4),
        "total_traces": total_traces,
        "source": "Langfuse /api/public/metrics/daily",
    }


def fetch_langfuse_tokens(days: int = 7, http_get: Optional[HTTP_GET] = None) -> dict[str, Any]:
    url = "%s/api/public/metrics/daily?limit=%d" % (
        C.LANGFUSE_URL.rstrip("/"),
        max(1, days),
    )
    if http_get is not None:
        raw = http_get(url)
    else:
        raw = curl_get(url, auth=(C.LANGFUSE_PK, C.LANGFUSE_SK))
    return parse_langfuse_daily(raw, days=days)


def _parse_util_percent(text: str) -> float:
    m = re.search(r"(\d+(?:\.\d+)?)\s*%", text or "")
    return float(m.group(1)) if m else 0.0


def parse_nvidia_sample(raw: str) -> dict[str, Any]:
    line = (raw or "").strip().splitlines()[0] if (raw or "").strip() else ""
    parts = [p.strip() for p in line.split(",")]
    if len(parts) < 6:
        raise ValueError("unexpected nvidia-smi: %r" % raw)
    mem_used = int(parts[4].replace("MiB", "").strip() or "0")
    mem_total = int(parts[5].replace("MiB", "").strip() or "0")
    return {
        "name": parts[1],
        "temperature_c": int(parts[2]),
        "utilization_pct": _parse_util_percent(parts[3]),
        "memory_used_mib": mem_used,
        "memory_total_mib": mem_total,
        "memory_pct": round(100.0 * mem_used / mem_total, 1) if mem_total else 0.0,
    }


def sample_gpu_now(runner: Optional[HTTP_GET] = None) -> dict[str, Any]:
    cmd = (
        "nvidia-smi --query-gpu=index,name,temperature.gpu,utilization.gpu,"
        "memory.used,memory.total --format=csv,noheader"
    )
    gpu = parse_nvidia_sample(gpu_exec(cmd, runner=runner))
    gpu["host"] = C.GPU_HOST
    gpu["sampled_at"] = datetime.now(TZ8).isoformat()
    return gpu


def update_gpu_history(sample: dict[str, Any]) -> Path:
    today = datetime.now(TZ8).date().isoformat()
    path = data_dir() / "gpu_daily.json"
    hist: dict[str, Any] = {}
    if path.is_file():
        hist = json.loads(path.read_text(encoding="utf-8"))
    day = hist.get(today) or {"samples": [], "util_max": 0.0, "util_sum": 0.0}
    util = float(sample.get("utilization_pct") or 0)
    day["samples"].append(sample)
    day["util_max"] = max(float(day.get("util_max") or 0), util)
    day["util_sum"] = float(day.get("util_sum") or 0) + util
    day["util_avg"] = round(day["util_sum"] / len(day["samples"]), 2)
    day["memory_pct_last"] = sample.get("memory_pct")
    day["gpu_name"] = sample.get("name")
    hist[today] = day
    cutoff = (datetime.now(TZ8).date() - timedelta(days=int(C.HOMEWORK6_DAYS) + 1)).isoformat()
    hist = {k: v for k, v in hist.items() if k >= cutoff}
    path.write_text(json.dumps(hist, ensure_ascii=False, indent=2), encoding="utf-8")
    return path


def gpu_weekly_from_history(days: int = 7) -> dict[str, Any]:
    path = data_dir() / "gpu_daily.json"
    hist: dict[str, Any] = {}
    if path.is_file():
        hist = json.loads(path.read_text(encoding="utf-8"))
    keys = sorted(hist.keys())[-days:]
    daily = []
    utils = []
    for k in keys:
        row = hist[k]
        avg = float(row.get("util_avg") or 0)
        utils.append(avg)
        daily.append(
            {
                "date": k,
                "util_avg_pct": avg,
                "util_max_pct": float(row.get("util_max") or 0),
                "memory_pct": row.get("memory_pct_last"),
                "gpu_name": row.get("gpu_name"),
            }
        )
    return {
        "source": "local gpu_daily.json (SSH nvidia-smi samples)",
        "days": days,
        "daily": daily,
        "avg_utilization_pct": round(sum(utils) / len(utils), 2) if utils else None,
        "max_utilization_pct": round(max(utils), 2) if utils else None,
        "note": "本环境暂无独立 Prometheus TSDB，脚本每次运行采样并滚动保留 7 日",
    }


def parse_prometheus_query_range(raw: str) -> Optional[float]:
    data = json.loads(raw)
    if data.get("status") != "success":
        return None
    result = (data.get("data") or {}).get("result") or []
    if not result:
        return None
    values = result[0].get("values") or []
    nums = []
    for _ts, val in values:
        try:
            nums.append(float(val))
        except (TypeError, ValueError):
            continue
    if not nums:
        return None
    return round(sum(nums) / len(nums), 2)


def fetch_prometheus_gpu_avg(days: int = 7, http_get: Optional[HTTP_GET] = None) -> Optional[dict[str, Any]]:
    base = (getattr(C, "HOMEWORK6_PROMETHEUS_URL", "") or "").strip().rstrip("/")
    if not base:
        return None
    end = int(datetime.now(timezone.utc).timestamp())
    start = end - days * 86400
    for query in PROM_QUERIES:
        params = urllib.parse.urlencode(
            {"query": query, "start": start, "end": end, "step": "3600"}
        )
        url = "%s/api/v1/query_range?%s" % (base, params)
        try:
            raw = http_get(url) if http_get else curl_get(url)
            avg = parse_prometheus_query_range(raw)
            if avg is not None:
                return {
                    "source": "Prometheus query_range",
                    "prometheus_url": base,
                    "query": query,
                    "avg_utilization_pct": avg,
                    "days": days,
                }
        except Exception:
            continue
    return None


def fetch_gpu_weekly(days: int = 7, runner: Optional[HTTP_GET] = None) -> dict[str, Any]:
    prom = fetch_prometheus_gpu_avg(days=days, http_get=runner)
    sample = sample_gpu_now(runner=runner if runner else None)
    update_gpu_history(sample)
    local = gpu_weekly_from_history(days=days)
    out = {
        "host": C.GPU_HOST,
        "current": sample,
        "weekly": local,
        "prometheus": prom,
    }
    if prom:
        out["avg_utilization_pct"] = prom["avg_utilization_pct"]
        out["source"] = prom["source"]
    else:
        out["avg_utilization_pct"] = local.get("avg_utilization_pct")
        out["source"] = local.get("source")
    return out


def build_facts(token_stats: dict[str, Any], gpu_stats: dict[str, Any]) -> dict[str, Any]:
    start = (datetime.now(TZ8).date() - timedelta(days=int(C.HOMEWORK6_DAYS) - 1)).isoformat()
    end = datetime.now(TZ8).date().isoformat()
    return {
        "period": {"start": start, "end": end, "days": C.HOMEWORK6_DAYS},
        "langfuse": token_stats,
        "gpu": gpu_stats,
    }


def qwen_weekly_report(facts: dict[str, Any], http_post: Optional[Callable[..., str]] = None) -> str:
    system = (
        "你是企业 AI 平台运维分析师。根据给定的 JSON 事实数据，写一份简洁的中文周报 Markdown。"
        "必须包含：概览、Token 消耗趋势、GPU 利用情况、风险与建议。"
        "只使用 JSON 里有的数字，不要编造。直接输出最终周报正文，不要输出思考过程。"
    )
    user = "请基于以下事实生成周报：\n" + json.dumps(facts, ensure_ascii=False, indent=2)
    body = {
        "model": C.QWEN_MODEL,
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        "temperature": 0.3,
        "max_tokens": 2048,
    }
    if http_post is not None:
        raw = http_post(json.dumps(body, ensure_ascii=False))
    else:
        url = C.VLLM_BASE.rstrip("/") + "/chat/completions"
        req = urllib.request.Request(
            url,
            data=json.dumps(body).encode("utf-8"),
            headers={
                "Content-Type": "application/json",
                "Authorization": "Bearer %s" % C.VLLM_KEY,
            },
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=180) as resp:
            raw = resp.read().decode("utf-8", "replace")
    data = json.loads(raw)
    msg = data["choices"][0]["message"]
    content = (msg.get("content") or msg.get("reasoning") or "").strip()
    if not content:
        raise RuntimeError("Qwen 返回空内容: %s" % raw[:600])
    return content


def send_dingtalk_markdown(title: str, text: str, http_post: Optional[Callable[..., str]] = None) -> dict[str, Any]:
    kw = getattr(C, "HOMEWORK6_DINGTALK_KEYWORD", "") or ""
    # 钉钉只校验 markdown.text 正文是否含关键词，标题不算
    if kw and kw not in text:
        text = "**%s**\n\n%s" % (kw, text)
    if kw and kw not in title:
        title = "%s · %s" % (kw, title)
    body = {"msgtype": "markdown", "markdown": {"title": title, "text": text}}
    if http_post is not None:
        raw = http_post(json.dumps(body, ensure_ascii=False))
        return json.loads(raw)
    return curl_post_json(C.HOMEWORK6_DINGTALK_WEBHOOK, body)


def save_report(text: str, facts: dict[str, Any]) -> Path:
    ts = datetime.now(TZ8).strftime("%Y%m%d_%H%M%S")
    md = data_dir() / ("weekly_report_%s.md" % ts)
    meta = data_dir() / ("weekly_facts_%s.json" % ts)
    md.write_text(text, encoding="utf-8")
    meta.write_text(json.dumps(facts, ensure_ascii=False, indent=2), encoding="utf-8")
    return md


def run_weekly(dry_run: bool = False) -> dict[str, Any]:
    token_stats = fetch_langfuse_tokens(days=int(C.HOMEWORK6_DAYS))
    gpu_stats = fetch_gpu_weekly(days=int(C.HOMEWORK6_DAYS))
    facts = build_facts(token_stats, gpu_stats)
    report = qwen_weekly_report(facts)
    path = save_report(report, facts)
    result: dict[str, Any] = {
        "report_path": str(path),
        "report_preview": report[:800],
        "total_tokens": token_stats.get("total_tokens"),
        "avg_gpu_util": gpu_stats.get("avg_utilization_pct"),
    }
    if dry_run:
        result["dingtalk"] = {"skipped": True, "reason": "dry_run"}
        return result
    title = getattr(C, "HOMEWORK6_REPORT_TITLE", "AI 平台运维周报")
    ding = send_dingtalk_markdown(title, report)
    result["dingtalk"] = ding
    return result