#!/usr/bin/env python3
"""
End-to-end Test für ComfyUI Image Gen API.
Iteriert über model × quality × style, generiert Bilder, schreibt Report.
Usage: python3 scripts/comfy_e2e_test.py [--api-url URL] [--output-dir DIR]
"""
import os, time, json, base64, pathlib, argparse


def parse_args():
    p = argparse.ArgumentParser(description="ComfyUI E2E image generation test")
    p.add_argument("--api-url", default="http://localhost:8000/v1/images/generations")
    p.add_argument("--output-dir", default=str(pathlib.Path.home() / ".hermes/cron/output/comfy_e2e"))
    p.add_argument("--models", default="ideogram4,flux2", help="Comma-separated models to test")
    p.add_argument("--qualities", default="low,medium,high", help="Comma-separated qualities")
    p.add_argument("--styles", default=",photorealistic", help="Comma-separated styles (empty for none)")
    p.add_argument("--size", default="1024x1024")
    p.add_argument("--n", type=int, default=1)
    p.add_argument("--timeout", type=int, default=1800, help="HTTP timeout per request in seconds")
    return p.parse_args()


def generate_image(api_url: str, payload: dict, timeout: int):
    import urllib.request
    req = urllib.request.Request(
        api_url,
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
        method="POST"
    )
    start = time.time()
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            data = json.loads(resp.read().decode())
            elapsed = time.time() - start
            b64 = data["data"][0]["b64_json"]
            return base64.b64decode(b64), elapsed
    except Exception as e:
        return None, str(e)


def main():
    args = parse_args()
    out_dir = pathlib.Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    models = [m.strip() for m in args.models.split(",")]
    qualities = [q.strip() for q in args.qualities.split(",")]
    styles = [s.strip() for s in args.styles.split(",")]

    PROMPTS = {
        "ideogram4": "A futuristic city skyline at sunset, neon lights, detailed architecture",
        "flux2": "A serene mountain lake reflection, golden hour lighting, highly detailed",
        "flux1-dev": "A cyberpunk street market at night, volumetric fog, neon signs",
        "flux1-schnell": "A minimalist abstract painting, bold colors, geometric shapes",
    }

    RESULTS = []
    total = len(models) * len(qualities) * len(styles)
    print(f"Output: {out_dir}")
    print(f"Total: {total} combinations")
    print("=" * 60)

    for model in models:
        for quality in qualities:
            for style in styles:
                tag = f"{model}_{quality}_{style or 'default'}"
                prompt = PROMPTS.get(model, "A beautiful landscape")
                full_prompt = f"{style}, {prompt}".strip(", ") if style else prompt
                payload = {
                    "model": model,
                    "prompt": full_prompt,
                    "n": args.n,
                    "size": args.size,
                    "quality": quality,
                    "response_format": "b64_json"
                }
                print(f"\n[{tag}] {full_prompt[:60]}...")
                img_bytes, info = generate_image(args.api_url, payload, args.timeout)
                if img_bytes:
                    fname = f"{tag}_{int(time.time())}.png"
                    fpath = out_dir / fname
                    fpath.write_bytes(img_bytes)
                    print(f"  ✅ {info:.1f}s, {len(img_bytes)/1024:.0f} KB → {fname}")
                    RESULTS.append({
                        "model": model, "quality": quality, "style": style,
                        "status": "OK", "time_s": round(info, 1),
                        "size_kb": round(len(img_bytes)/1024, 1), "file": str(fpath)
                    })
                else:
                    print(f"  ❌ {info}")
                    RESULTS.append({
                        "model": model, "quality": quality, "style": style,
                        "status": "FAILED", "error": str(info)
                    })

    # JSON report
    report = out_dir / "report.json"
    report.write_text(json.dumps(RESULTS, indent=2))

    # Markdown report
    md = out_dir / "report.md"
    lines = [
        "# ComfyUI E2E Test Report",
        f"\n**Zeit:** {time.strftime('%Y-%m-%d %H:%M:%S')}",
    ]
    ok = sum(1 for r in RESULTS if r["status"] == "OK")
    fail = sum(1 for r in RESULTS if r["status"] == "FAILED")
    lines.append(f"**Ergebnis:** {ok} OK / {fail} FAILED\n")
    lines.append("| Modell | Quality | Style | Status | Zeit | Datei |")
    lines.append("|---|---|---|---|---|---|")
    for r in RESULTS:
        style = r.get("style", "") or "—"
        status = "✅" if r["status"] == "OK" else "❌"
        t = f"{r.get('time_s', '—')}s" if "time_s" in r else "—"
        f = r.get("file", "").split("/")[-1] if "file" in r else "—"
        lines.append(f"| {r['model']} | {r['quality']} | {style} | {status} {r['status']} | {t} | {f} |")
    md.write_text("\n".join(lines))
    print(f"\n{'=' * 60}")
    print(f"Reports: {report} , {md}")
    print(f"Bilder: {out_dir}")


if __name__ == "__main__":
    main()