#!/usr/bin/env python3
"""Push smoke test — create a fresh unique repo, build commits, push, report timing.

Usage:
    python tools/smoke_push.py xs          # 1 commit,   5 files,  512 B/file
    python tools/smoke_push.py s           # 10 commits, 10 files, 1 KB/file
    python tools/smoke_push.py m           # 100 commits,20 files, 2 KB/file
    python tools/smoke_push.py l           # 500 commits,50 files, 4 KB/file
    python tools/smoke_push.py xl          # 1000 commits,100 files,8 KB/file
    python tools/smoke_push.py xs s m      # run multiple sizes in sequence
    python tools/smoke_push.py --hub https://staging.musehub.ai xs
"""
from __future__ import annotations

import json
import random
import shutil
import string
import subprocess
import sys
import time
from pathlib import Path

HUB_REPO_CTX = str(Path(__file__).parent.parent)  # musehub repo root

SIZES = {
    "xs": (1,    5,    512),
    "s":  (10,   10,   1_024),
    "m":  (100,  20,   2_048),
    "l":  (500,  50,   4_096),
    "xl": (1000, 100,  8_192),
}


def run(cmd: list[str], cwd: str | None = None, check: bool = True) -> subprocess.CompletedProcess:
    return subprocess.run(cmd, cwd=cwd, capture_output=True, text=True, check=check)


def rnd_content(size: int, seed: int) -> str:
    rng = random.Random(seed)
    return "".join(random.Random(seed).choice(string.ascii_letters + string.digits + "\n")
                   for _ in range(size))


def smoke(label: str, hub: str) -> dict:
    n_commits, n_files, file_size = SIZES[label]
    ts = int(time.time())
    slug = f"smoke-{label}-{ts}"
    repo_path = Path(f"/tmp/{slug}")

    total_kb = n_commits * n_files * file_size // 1024
    print(f"\n{'='*64}")
    print(f"  {label.upper()} — {n_commits} commits × {n_files} files × {file_size}B ≈ {total_kb} KB")
    print(f"  repo: gabriel/{slug}  hub: {hub}")
    print(f"{'='*64}")

    # Fresh local repo
    repo_path.mkdir(parents=True, exist_ok=True)
    run(["muse", "init"], cwd=str(repo_path))
    run(["muse", "remote", "set-url", "local", f"{hub}/gabriel/{slug}"], cwd=str(repo_path))

    # Create hub repo
    print(f"  creating hub repo…", flush=True)
    r = run(["muse", "-C", HUB_REPO_CTX, "hub", "repo", "create",
             "--name", slug, "--visibility", "public", "--hub", hub], check=False)
    if r.returncode != 0:
        print(f"  ❌ hub create failed: {r.stderr[:200]}")
        shutil.rmtree(repo_path, ignore_errors=True)
        return {"label": label, "error": "hub create failed"}

    # Build commits
    print(f"  building {n_commits} commit(s)…", flush=True)
    t_build = time.monotonic()
    src_dir = repo_path / "src"
    src_dir.mkdir()
    for ci in range(n_commits):
        for fi in range(n_files):
            (src_dir / f"file_{fi:04d}.txt").write_text(rnd_content(file_size, ci * 10000 + fi))
        run(["muse", "code", "add", "."], cwd=str(repo_path))
        run(["muse", "commit", "-m", f"c{ci+1}",
             "--agent-id", "smoke", "--model-id", "none", "--sign"], cwd=str(repo_path))
        if n_commits >= 100 and (ci + 1) % 100 == 0:
            print(f"    {ci+1}/{n_commits}", flush=True)
    print(f"  built in {time.monotonic()-t_build:.1f}s", flush=True)

    # Object count from HEAD manifest
    try:
        m = json.loads(run(["muse", "read", "--json", "--manifest"], cwd=str(repo_path)).stdout)
        n_objects = len(m.get("manifest", {}))
    except Exception:
        n_objects = -1
    print(f"  HEAD snapshot: {n_objects} unique objects")

    # First push
    print(f"  pushing…", flush=True)
    t0 = time.monotonic()
    p = run(["muse", "push", "local", "main"], cwd=str(repo_path), check=False)
    elapsed = time.monotonic() - t0

    if p.returncode != 0:
        print(f"  ❌ FAILED in {elapsed:.3f}s")
        print(p.stderr[-400:])
        shutil.rmtree(repo_path, ignore_errors=True)
        return {"label": label, "n_commits": n_commits, "n_objects": n_objects,
                "error": p.stderr[-120:].strip(), "first_push_s": round(elapsed, 3)}

    print(f"  ✅ first push: {elapsed:.3f}s")
    # Show key server timing lines
    for line in p.stdout.splitlines():
        if any(k in line for k in ["TOTAL server", "stored", "✅ Pushed"]):
            print(f"     {line.strip()}")

    # Re-push (all objects already in DB — tests O(1) presign path)
    print(f"  re-pushing (already-stored)…", flush=True)
    t1 = time.monotonic()
    p2 = run(["muse", "push", "local", "main"], cwd=str(repo_path), check=False)
    elapsed2 = time.monotonic() - t1
    print(f"  {'✅' if p2.returncode == 0 else '❌'} re-push: {elapsed2:.3f}s")

    shutil.rmtree(repo_path, ignore_errors=True)
    return {
        "label": label, "n_commits": n_commits, "n_objects": n_objects,
        "first_push_s": round(elapsed, 3), "repush_s": round(elapsed2, 3), "ok": True,
    }


def main() -> None:
    args = sys.argv[1:]
    hub = "https://localhost:1337"

    # --hub flag
    if "--hub" in args:
        i = args.index("--hub")
        hub = args[i + 1]
        args = args[:i] + args[i + 2:]

    if not args:
        print(__doc__)
        sys.exit(1)

    labels = [a.lower() for a in args]
    for lbl in labels:
        if lbl not in SIZES:
            print(f"unknown size '{lbl}' — choose from: {', '.join(SIZES)}")
            sys.exit(1)

    results = [smoke(lbl, hub) for lbl in labels]

    if len(results) > 1:
        print(f"\n{'='*64}")
        print(f"  SUMMARY")
        print(f"{'='*64}")
        print(f"{'Size':<6} {'Commits':>8} {'Objects':>8} {'First':>10} {'Re-push':>9}  Status")
        print("-" * 50)
        for r in results:
            if "error" in r and "first_push_s" not in r:
                print(f"{r['label'].upper():<6} {r.get('n_commits','?'):>8} {r.get('n_objects','?'):>8} "
                      f"{'FAILED':>10} {'—':>9}  ❌")
            else:
                print(f"{r['label'].upper():<6} {r['n_commits']:>8} {r['n_objects']:>8} "
                      f"{r['first_push_s']:>9.3f}s {r.get('repush_s',0):>8.3f}s  ✅")


if __name__ == "__main__":
    main()