#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
devperf_bench.py
Mac Studio (M2 Ultra) 开发链路网络性能基准测试 + 专业UI
- 统一代理：127.0.0.1:7890（ClashX）
- 任务：brew / mvn / docker / kubectl / helm / pip / git
- 指标：latency / bytes / throughput / success
- UI：Chart.js + Tailwind，一键可视化
"""

import asyncio, os, sys, json, time, shutil, tempfile, hashlib, psutil, pathlib
from typing import List, Dict, Any, Optional
from fastapi import FastAPI, BackgroundTasks, Query
from fastapi.responses import HTMLResponse, JSONResponse, PlainTextResponse

# =========================
# 全局配置（可自行调整）
# =========================
PROXY_HTTP = os.environ.get("http_proxy", "http://127.0.0.1:7890")
PROXY_HTTPS = os.environ.get("https_proxy", "http://127.0.0.1:7890")

# 测试对象（使用更大型的包/镜像/仓库进行专业测试）
CFG = {
    "concurrency": 4,          # 并发度
    "repeat": 1,               # 每类任务重复次数
    "workdir": ".devperf_tmp", # 工作目录
    "targets": {
        "brew": {
            "desc": "brew 拉取大型 bottle（用下载替代）",
            # 使用 brew 下载一个较大型的包（不会安装到系统，尽量不污染环境）
            "cmd": "brew fetch --force --retry --quiet --deps node python docker",
            "size_probe": "du -sb $(brew --cache) | head -n 1 | cut -f1"
        },
        "mvn": {
            "desc": "maven 预拉大型项目依赖（Spring Boot Web Starter）",
            "cmd": "mvn -q -Dmaven.repo.local=.m2_repo -f pom.xml dependency:go-offline",
            "bootstrap": {
                "pom": """<project xmlns="http://maven.apache.org/POM/4.0.0" 
  xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" 
  xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 
  http://maven.apache.org/xsd/maven-4.0.0.xsd">
  <modelVersion>4.0.0</modelVersion>
  <groupId>bench</groupId><artifactId>bench</artifactId><version>1.0</version>
  <packaging>jar</packaging>
  <properties>
    <maven.compiler.source>17</maven.compiler.source>
    <maven.compiler.target>17</maven.compiler.target>
  </properties>
  <dependencies>
    <dependency>
      <groupId>org.springframework.boot</groupId>
      <artifactId>spring-boot-starter-web</artifactId>
      <version>3.2.0</version>
    </dependency>
    <dependency>
      <groupId>org.springframework.boot</groupId>
      <artifactId>spring-boot-starter-data-jpa</artifactId>
      <version>3.2.0</version>
    </dependency>
    <dependency>
      <groupId>com.fasterxml.jackson.core</groupId>
      <artifactId>jackson-databind</artifactId>
      <version>2.15.2</version>
    </dependency>
    <dependency>
      <groupId>org.hibernate</groupId>
      <artifactId>hibernate-core</artifactId>
      <version>6.3.1.Final</version>
    </dependency>
  </dependencies>
</project>"""
            },
            "size_dir": ".m2_repo"  # 以目录体积估算
        },
        "docker": {
            "desc": "docker 拉取大型镜像（Node.js + Python + Java 环境）",
            "images": ["node:21-alpine", "python:3.12-slim", "openjdk:17-jdk-slim", "postgres:16-alpine"],
            "cmd_tpl": "docker pull {image}"
        },
        "kubectl": {
            "desc": "k8s 组件镜像（用 docker pull 官方镜像代替触达）",
            "images": [
                "registry.k8s.io/kube-apiserver:v1.28.0",
                "registry.k8s.io/kube-controller-manager:v1.28.0",
                "registry.k8s.io/kube-scheduler:v1.28.0",
                "registry.k8s.io/coredns/coredns:v1.10.1"
            ],
            "cmd_tpl": "docker pull {image}"
        },
        "helm": {
            "desc": "helm 拉取大型 chart（Kubernetes Dashboard + Prometheus）",
            "init": [
                "helm repo add stable https://charts.helm.sh/stable",
                "helm repo add prometheus-community https://prometheus-community.github.io/helm-charts",
                "helm repo update"
            ],
            "pull": "helm pull prometheus-community/kube-prometheus --destination . && helm pull stable/kubernetes-dashboard --destination ."
        },
        "pip": {
            "desc": "pip 下载大型库（Django, FastAPI, Pandas等）",
            "cmd": "pip download --no-deps --dest . 'django>=4.2' fastapi uvicorn pandas numpy requests pydantic",
            "size_glob": "*.whl"
        },
        "git": {
            "desc": "GitHub 克隆大型开源项目（浅克隆）",
            "repos": [
                "https://github.com/microsoft/vscode.git",
                "https://github.com/torvalds/linux.git",
                "https://github.com/facebook/react.git",
                "https://github.com/angular/angular.git"
            ],
            "cmd_tpl": "git clone --depth 1 --single-branch {repo}"
        }
    }
}

RESULTS: List[Dict[str, Any]] = []
RUNNING = False

# ============ 工具函数 ============

def human(n: float) -> str:
    for u in ["B","KB","MB","GB","TB"]:
        if n < 1024: return f"{n:.2f} {u}"
        n /= 1024
    return f"{n:.2f} PB"

async def run_cmd(cmd: str, cwd: Optional[str]=None, env: Optional[dict]=None, timeout: int=0) -> Dict[str, Any]:
    """运行命令并返回 {rc, out, err, elapsed}"""
    start = time.time()
    try:
        proc = await asyncio.create_subprocess_shell(
            cmd,
            cwd=cwd,
            env=env or os.environ.copy(),
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        out, err = await asyncio.wait_for(proc.communicate(), timeout=timeout or None)
        elapsed = time.time() - start
        return {"rc": proc.returncode, "out": out.decode(errors="ignore"), "err": err.decode(errors="ignore"), "elapsed": elapsed}
    except asyncio.TimeoutError:
        return {"rc": -1, "out": "", "err": f"timeout after {timeout}s", "elapsed": time.time() - start}

def dir_size_bytes(path: str) -> int:
    total = 0
    p = pathlib.Path(path)
    if not p.exists(): return 0
    for f in p.rglob("*"):
        if f.is_file():
            try: total += f.stat().st_size
            except: pass
    return total

def glob_total_bytes(glob_pattern: str, cwd: str) -> int:
    total = 0
    for f in pathlib.Path(cwd).glob(glob_pattern):
        if f.is_file():
            try: total += f.stat().st_size
            except: pass
    return total

def record(task: str, label: str, bytes_n: int, elapsed: float, rc: int, extra: Dict[str, Any]=None):
    thr = (bytes_n/elapsed/1024/1024) if (bytes_n>0 and elapsed>0) else 0.0
    RESULTS.append({
        "ts": time.strftime("%Y-%m-%d %H:%M:%S"),
        "task": task,
        "label": label,
        "bytes": bytes_n,
        "elapsed_sec": round(elapsed, 3),
        "throughput_mbps": round(thr, 2),
        "success": (rc==0),
        **(extra or {})
    })

# ============ 每类任务的执行器 ============

async def bench_brew(base: str):
    td = pathlib.Path(base, "brew"); td.mkdir(parents=True, exist_ok=True)
    # 清理缓存项以避免复用（可选）
    # await run_cmd("brew cleanup --prune=all", cwd=str(td))
    res = await run_cmd(CFG["targets"]["brew"]["cmd"], cwd=str(td))
    size_res = await run_cmd(CFG["targets"]["brew"]["size_probe"], cwd=str(td))
    bytes_n = int(size_res["out"].strip()) if size_res["rc"]==0 and size_res["out"].strip().isdigit() else 0
    record("brew", "node+python+docker", bytes_n, res["elapsed"], res["rc"], {"stderr": res["err"][-200:]})

async def bench_mvn(base: str):
    td = pathlib.Path(base, "mvn"); td.mkdir(parents=True, exist_ok=True)
    (td / "pom.xml").write_text(CFG["targets"]["mvn"]["bootstrap"]["pom"], encoding="utf-8")
    # 先清目录
    shutil.rmtree(td / ".m2_repo", ignore_errors=True)
    res = await run_cmd(CFG["targets"]["mvn"]["cmd"], cwd=str(td))
    bytes_n = dir_size_bytes(str(td / ".m2_repo"))
    record("mvn", "spring-boot-starter-web", bytes_n, res["elapsed"], res["rc"], {"stderr": res["err"][-200:]})

async def bench_docker(base: str):
    td = pathlib.Path(base, "docker"); td.mkdir(parents=True, exist_ok=True)
    for img in CFG["targets"]["docker"]["images"]:
        res = await run_cmd(CFG["targets"]["docker"]["cmd_tpl"].format(image=img), cwd=str(td))
        # 用 inspect 拿镜像大小
        insp = await run_cmd(f"docker image inspect {img} --format='{{{{.Size}}}}'")
        bytes_n = int(insp["out"].strip()) if insp["rc"]==0 and insp["out"].strip().isdigit() else 0
        record("docker", img, bytes_n, res["elapsed"], res["rc"], {"stderr": res["err"][-200:]})

async def bench_kubectl(base: str):
    td = pathlib.Path(base, "kubectl"); td.mkdir(parents=True, exist_ok=True)
    for img in CFG["targets"]["kubectl"]["images"]:
        res = await run_cmd(CFG["targets"]["kubectl"]["cmd_tpl"].format(image=img), cwd=str(td))
        insp = await run_cmd(f"docker image inspect {img} --format='{{{{.Size}}}}'")
        bytes_n = int(insp["out"].strip()) if insp["rc"]==0 and insp["out"].strip().isdigit() else 0
        record("k8s-image", img, bytes_n, res["elapsed"], res["rc"], {"stderr": res["err"][-200:]})

async def bench_helm(base: str):
    td = pathlib.Path(base, "helm"); td.mkdir(parents=True, exist_ok=True)
    # 初始化仓库
    for cmd in CFG["targets"]["helm"]["init"]:
        await run_cmd(cmd, cwd=str(td))
    # 删除旧文件，避免缓存干扰
    for f in pathlib.Path(td).glob("*.tgz"): f.unlink(missing_ok=True)
    res = await run_cmd(CFG["targets"]["helm"]["pull"], cwd=str(td))
    bytes_n = glob_total_bytes("*.tgz", str(td))
    record("helm", "prometheus+k8s-dashboard", bytes_n, res["elapsed"], res["rc"], {"stderr": res["err"][-200:]})

async def bench_pip(base: str):
    td = pathlib.Path(base, "pip"); td.mkdir(parents=True, exist_ok=True)
    # 清空旧轮子
    for f in pathlib.Path(td).glob("*.whl"): f.unlink(missing_ok=True)
    res = await run_cmd(CFG["targets"]["pip"]["cmd"], cwd=str(td))
    bytes_n = glob_total_bytes("*.whl", str(td))
    record("pip", "django+fastapi+pandas", bytes_n, res["elapsed"], res["rc"], {"stderr": res["err"][-200:]})

async def bench_git(base: str):
    td = pathlib.Path(base, "git"); td.mkdir(parents=True, exist_ok=True)
    for repo in CFG["targets"]["git"]["repos"]:
        name = repo.rsplit("/",1)[-1].replace(".git", "")
        dst = pathlib.Path(td, name)
        if dst.exists(): shutil.rmtree(dst, ignore_errors=True)
        res = await run_cmd(CFG["targets"]["git"]["cmd_tpl"].format(repo=repo), cwd=str(td))
        bytes_n = dir_size_bytes(str(dst))
        record("git", name, bytes_n, res["elapsed"], res["rc"], {"stderr": res["err"][-200:]})

TASK_FUNCS = {
    "brew": bench_brew,
    "mvn": bench_mvn,
    "docker": bench_docker,
    "kubectl": bench_kubectl,
    "helm": bench_helm,
    "pip": bench_pip,
    "git": bench_git
}

# ============ 调度器 ============

async def run_all(concurrency: int, repeat: int):
    global RUNNING
    if RUNNING: return
    RUNNING = True
    RESULTS.clear()
    base = CFG["workdir"]
    shutil.rmtree(base, ignore_errors=True)
    pathlib.Path(base).mkdir(parents=True, exist_ok=True)

    # 确保代理在环境变量中
    os.environ["http_proxy"] = PROXY_HTTP
    os.environ["https_proxy"] = PROXY_HTTPS

    # 构建任务队列
    jobs = []
    for _ in range(repeat):
        for key in TASK_FUNCS.keys():
            jobs.append((key, TASK_FUNCS[key]))

    sem = asyncio.Semaphore(concurrency)

    async def wrapper(name, coro):
        async with sem:
            try:
                await coro(CFG["workdir"])
            except Exception as e:
                record(name, "exception", 0, 0.0, -1, {"exception": str(e)})

    await asyncio.gather(*[wrapper(n, f) for n,f in jobs])

    # 写入结果
    pathlib.Path(base, "results.json").write_text(json.dumps(RESULTS, ensure_ascii=False, indent=2), encoding="utf-8")
    RUNNING = False

# ============ FastAPI + UI ============

app = FastAPI(title="DevPerf Bench (Mac Studio M2 Ultra)")

@app.get("/", response_class=HTMLResponse)
def index():
    return HTMLResponse("""
<!DOCTYPE html><html><head>
<meta charset="utf-8"/>
<meta name="viewport" content="width=device-width, initial-scale=1"/>
<title>DevPerf Bench · M2 Ultra</title>
<script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
<script>
async function fetchJSON(url){const r=await fetch(url);return await r.json()}
async function runBench(){
  const cc=document.getElementById("concurrency").value;
  const rp=document.getElementById("repeat").value;
  document.getElementById("status").innerText="Running..."
  await fetch(`/run?concurrency=${cc}&repeat=${rp}`); 
  await refresh()
  document.getElementById("status").innerText="Done"
}
async function refresh(){
  const data=await fetchJSON('/results')
  // 聚合
  const byTask = {}
  data.forEach(d=>{
    if(!byTask[d.task]) byTask[d.task]=[]
    byTask[d.task].push(d)
  })
  // 表格
  const tbody=document.getElementById("tbody")
  tbody.innerHTML=""
  data.forEach(d=>{
    const tr=document.createElement("tr");
    tr.innerHTML=`<td class="px-2 py-1">${d.ts}</td>
      <td class="px-2 py-1">${d.task}</td>
      <td class="px-2 py-1">${d.label}</td>
      <td class="px-2 py-1">${(d.bytes/1024/1024).toFixed(2)} MB</td>
      <td class="px-2 py-1">${d.elapsed_sec}s</td>
      <td class="px-2 py-1">${d.throughput_mbps} MB/s</td>
      <td class="px-2 py-1 ${d.success?'text-green-600':'text-red-600'}">${d.success?'OK':'FAIL'}</td>`
    tbody.appendChild(tr)
  })
  // 图表：每类任务的吞吐
  const labels=[], thr=[]
  Object.keys(byTask).forEach(k=>{
    const arr=byTask[k]
    const avg = arr.reduce((a,b)=>a+b.throughput_mbps,0)/arr.length
    labels.push(k); thr.push(avg.toFixed(2))
  })
  const ctx=document.getElementById('chart').getContext('2d');
  if(window._chart) window._chart.destroy();
  window._chart=new Chart(ctx,{
    type:'bar',
    data:{labels:labels, datasets:[{label:'Avg Throughput (MB/s)', data:thr}]},
    options:{responsive:true, plugins:{legend:{display:true}}, scales:{y:{beginAtZero:true}}}
  })
}
document.addEventListener('DOMContentLoaded',()=>{refresh()})
</script>
<style>
:root { color-scheme: light dark; }
body { font-family: ui-sans-serif,system-ui,-apple-system,"Segoe UI",Roboto,"Helvetica Neue",Arial; margin: 0; padding: 1rem 1.2rem;}
.card { border: 1px solid #e5e7eb22; border-radius:14px; padding: 1rem; box-shadow: 0 1px 6px rgba(0,0,0,.06); margin-bottom: 1rem;}
.btn { padding: .6rem 1rem; border-radius: 10px; border: 1px solid #e5e7eb44; cursor: pointer;}
.input { padding: .5rem .7rem; border-radius: 10px; border: 1px solid #e5e7eb55; width: 6rem;}
table { width:100%; border-collapse: collapse; }
th, td { border-bottom: 1px dashed #e5e7eb33; font-size:.9rem;}
h1 { font-size: 1.4rem; margin:.2rem 0 .8rem}
h2 { font-size: 1.1rem; margin:.2rem 0 .4rem}
</style>
</head>
<body>
  <div class="card">
    <h1>DevPerf Bench · Mac Studio M2 Ultra</h1>
    <div style="display:flex; gap:.8rem; align-items:center; flex-wrap:wrap;">
      <div>并发度</div><input id="concurrency" class="input" type="number" min="1" value="4"/>
      <div>重复</div><input id="repeat" class="input" type="number" min="1" value="1"/>
      <button class="btn" onclick="runBench()">▶ 运行测试</button>
      <button class="btn" onclick="refresh()">⟲ 刷新</button>
      <div id="status" style="margin-left:auto; opacity:.8">Idle</div>
    </div>
  </div>

  <div class="card">
    <h2>吞吐对比（平均值）</h2>
    <canvas id="chart" height="120"></canvas>
  </div>

  <div class="card">
    <h2>明细</h2>
    <div style="overflow:auto; max-height:50vh;">
      <table>
        <thead><tr>
          <th class="px-2 py-1">时间</th>
          <th class="px-2 py-1">任务</th>
          <th class="px-2 py-1">标签</th>
          <th class="px-2 py-1">传输</th>
          <th class="px-2 py-1">时延</th>
          <th class="px-2 py-1">吞吐</th>
          <th class="px-2 py-1">状态</th>
        </tr></thead>
        <tbody id="tbody"></tbody>
      </table>
    </div>
  </div>

  <div style="opacity:.6; font-size:.85rem;">提示：结果已落盘 <code>.devperf_tmp/results.json</code>；你可改 CFG 里的镜像/仓库以贴合你的真实项目。</div>
</body></html>
    """)

@app.get("/run")
async def run(concurrency: int = Query(CFG["concurrency"], ge=1, le=64),
              repeat: int = Query(CFG["repeat"], ge=1, le=10),
              background_tasks: BackgroundTasks = None):
    """后台启动一次完整测试"""
    if RUNNING:
        return JSONResponse({"ok": False, "msg": "already running"})
    background_tasks.add_task(run_all, concurrency, repeat)
    return JSONResponse({"ok": True, "msg": "started"})

@app.get("/results")
def results():
    return JSONResponse(RESULTS)

@app.get("/raw")
def raw():
    """导出 JSON（便于你后续喂给 Grafana/ELK）"""
    return PlainTextResponse(json.dumps(RESULTS, ensure_ascii=False, indent=2))

if __name__ == "__main__":
    import uvicorn
    print(f"[i] Using proxy http={PROXY_HTTP} https={PROXY_HTTPS}")
    uvicorn.run(app, host="127.0.0.1", port=8000)