#!/usr/bin/env python3
"""
Stress test: run N concurrent end-to-end pipelines against the server.

Usage:
    export AIR_API_KEY="air_k1_..."
    python examples/stress_test.py [--clients 30] [--steps idea,literature,methods,results,paper]

Each client runs in its own thread with its own project. Logs per client are
written to stress_test_logs/<run_id>/<client_id>.log. A summary table is
printed at the end.
"""

import argparse
import os
import sys
import time
import threading
import uuid
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path

from dotenv import load_dotenv
import air

load_dotenv()

API_KEY = os.environ.get("AIR_API_KEY")
BASE_URL = os.environ.get("AIR_BASE_URL", "http://localhost:8000")

if not API_KEY:
    sys.exit("Set AIR_API_KEY first:  export AIR_API_KEY='air_k1_...'")

DATA_DESCRIPTION = "We want to do a basic study of simple harmonic oscillator."

_print_lock = threading.Lock()


def log(client_id: str, msg: str, file=None) -> None:
    ts = datetime.now().strftime("%H:%M:%S")
    line = f"[{ts}] [{client_id}] {msg}"
    if file:
        file.write(line + "\n")
        file.flush()
    with _print_lock:
        print(line, flush=True)


@dataclass
class ClientResult:
    client_id: str
    project_name: str
    steps_done: list[str] = field(default_factory=list)
    steps_failed: list[str] = field(default_factory=list)
    error: str | None = None
    duration: float = 0.0

    @property
    def success(self) -> bool:
        return self.error is None and not self.steps_failed


def run_client(client_id: str, steps: list[str], log_dir: Path) -> ClientResult:
    project_name = f"stress_{client_id}"
    result = ClientResult(client_id=client_id, project_name=project_name)
    log_path = log_dir / f"{client_id}.log"
    t0 = time.monotonic()

    with open(log_path, "w") as logfile:
        def out(msg: str) -> None:
            log(client_id, msg.strip(), logfile)

        out(f"START  project={project_name}")
        try:
            with air.AIR(api_key=API_KEY, base_url=BASE_URL) as client:
                project = client.create_project(
                    name=project_name,
                    data_description=DATA_DESCRIPTION,
                )
                out(f"Project created: {project}")

                if "idea" in steps:
                    out("[1/5] idea ...")
                    try:
                        idea = project.idea()
                        out(f"  idea: {len(idea)} chars")
                        result.steps_done.append("idea")
                    except Exception as e:
                        out(f"  idea FAILED: {e}")
                        result.steps_failed.append("idea")

                if "literature" in steps:
                    out("[2/5] literature ...")
                    try:
                        lit = project.literature()
                        out(f"  literature: {len(lit)} chars")
                        result.steps_done.append("literature")
                    except Exception as e:
                        out(f"  literature FAILED: {e}")
                        result.steps_failed.append("literature")

                if "methods" in steps:
                    out("[3/5] methods ...")
                    try:
                        meth = project.methods()
                        out(f"  methods: {len(meth)} chars")
                        result.steps_done.append("methods")
                    except Exception as e:
                        out(f"  methods FAILED: {e}")
                        result.steps_failed.append("methods")

                if "results" in steps:
                    out("[4/5] results (WebSocket + local execution) ...")
                    try:
                        res = project.results(
                            max_attempts=3,
                            max_steps=2,
                            on_output=lambda s: out(f"  ws| {s.strip()}"),
                        )
                        out(f"  results: {len(res)} chars")
                        result.steps_done.append("results")
                    except Exception as e:
                        out(f"  results FAILED: {e}")
                        result.steps_failed.append("results")

                if "paper" in steps and "results" not in steps:
                    # Upload a stub results.md so the paper step has something to work with.
                    fake_results = (
                        "## Results\n\n"
                        "This is a placeholder results section generated for stress testing.\n\n"
                        "### Experiment\n"
                        "A simple harmonic oscillator was simulated. "
                        "The system followed x(t) = A cos(ωt + φ) with ω = 1 rad/s, A = 1 m.\n\n"
                        "### Findings\n"
                        "- Period T = 2π s\n"
                        "- Energy conserved to within 0.01%\n"
                        "- Phase space trajectory forms a closed ellipse\n"
                    )
                    try:
                        project.write_file(f"Iteration0/input_files/results.md", fake_results)
                        out("  Uploaded fake results.md")
                    except Exception as e:
                        out(f"  WARNING: could not upload fake results.md: {e}")

                if "paper" in steps:
                    out("[5/5] paper ...")
                    try:
                        tex = project.paper(journal="NONE", add_citations=False)
                        out(f"  paper: {len(tex)} chars")
                        result.steps_done.append("paper")
                    except Exception as e:
                        out(f"  paper FAILED: {e}")
                        result.steps_failed.append("paper")

        except Exception as e:
            result.error = str(e)
            out(f"FATAL: {e}")

    result.duration = time.monotonic() - t0
    return result


def print_summary(results: list[ClientResult], total_duration: float) -> None:
    print("\n" + "=" * 72)
    print(f"STRESS TEST SUMMARY  ({len(results)} clients, {total_duration:.1f}s wall clock)")
    print("=" * 72)

    ok = [r for r in results if r.success]
    partial = [r for r in results if not r.success and not r.error]
    failed = [r for r in results if r.error]

    print(f"  Fully succeeded : {len(ok)}")
    print(f"  Partial (step ✗): {len(partial)}")
    print(f"  Fatal error     : {len(failed)}")
    print()

    header = f"{'CLIENT':<20} {'DURATION':>8}  {'DONE':<35} {'FAILED'}"
    print(header)
    print("-" * 72)
    for r in sorted(results, key=lambda r: r.client_id):
        done = ",".join(r.steps_done) or "-"
        fail = ",".join(r.steps_failed) or ("-" if not r.error else "FATAL")
        if r.error:
            fail = f"FATAL: {r.error[:30]}"
        print(f"  {r.client_id:<18} {r.duration:>7.1f}s  {done:<35} {fail}")
    print("=" * 72)


def main() -> None:
    parser = argparse.ArgumentParser(description="AIR server stress test")
    parser.add_argument("--clients", type=int, default=30, help="Number of concurrent clients")
    parser.add_argument(
        "--steps",
        default="idea,literature,methods,paper",
        help="Comma-separated pipeline steps to run. Omit 'results' to use a fake results.md.",
    )
    parser.add_argument("--max-workers", type=int, default=None,
                        help="Thread pool size (default: same as --clients)")
    args = parser.parse_args()

    steps = [s.strip() for s in args.steps.split(",")]
    n = args.clients
    workers = args.max_workers or n

    run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_dir = Path("stress_test_logs") / run_id
    log_dir.mkdir(parents=True, exist_ok=True)

    client_ids = [f"c{i:02d}_{uuid.uuid4().hex[:4]}" for i in range(n)]

    print(f"Stress test: {n} clients, steps={steps}, workers={workers}")
    print(f"Logs: {log_dir}/")
    print(f"Server: {BASE_URL}\n")

    t0 = time.monotonic()
    results: list[ClientResult] = []

    with ThreadPoolExecutor(max_workers=workers) as pool:
        futures = {pool.submit(run_client, cid, steps, log_dir): cid for cid in client_ids}
        for future in as_completed(futures):
            cid = futures[future]
            try:
                result = future.result()
            except Exception as e:
                result = ClientResult(client_id=cid, project_name=f"stress_{cid}", error=str(e))
            results.append(result)
            status = "OK" if result.success else ("PARTIAL" if not result.error else "FATAL")
            with _print_lock:
                print(f"  [{status}] {cid} finished in {result.duration:.1f}s  "
                      f"({len(results)}/{n} done)", flush=True)

    print_summary(results, time.monotonic() - t0)


if __name__ == "__main__":
    main()
