#!/usr/bin/env python3
"""ST-COUNT-01 — carga concurrente hacia count-api (vía gateway).

Equivalente práctico al k6 k6_count_batches.js para hosts sin k6 instalado.
Perfiles:
  shared8g  → ~1 req/s durante 60s, concurrencia máxima 3 (conservador)
  lab       → ~3 req/s durante 120s, concurrencia máxima 8

Usa PDFs reales de docs/plan_pruebas/data/ (rota entre ellos).

Ejemplos:
  python docs/plan_pruebas/scripts/st_count_01_bench.py
  python docs/plan_pruebas/scripts/st_count_01_bench.py --profile shared8g
  PROFILE=lab python docs/plan_pruebas/scripts/st_count_01_bench.py --duration 90
"""
from __future__ import annotations

import argparse
import csv
import json
import os
import statistics
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import date, datetime, timezone
from pathlib import Path

try:
    import requests
except ImportError:
    print("Requiere 'requests': pip install requests", file=sys.stderr)
    sys.exit(1)

ROOT = Path(__file__).resolve().parents[3]
DEFAULT_DATA = ROOT / "docs" / "plan_pruebas" / "data"
DEFAULT_INFORMES = ROOT / "docs" / "plan_pruebas" / "informes"

PROFILES = {
    "shared8g": {"rate": 1.0, "duration_s": 60, "max_workers": 3},
    "lab": {"rate": 3.0, "duration_s": 120, "max_workers": 8},
}


def list_pdfs(data_dir: Path) -> list[Path]:
    if not data_dir.is_dir():
        raise SystemExit(f"No existe DATA_DIR={data_dir}")
    pdfs = sorted({p.resolve() for p in list(data_dir.glob("*.pdf")) + list(data_dir.glob("*.PDF"))})
    if not pdfs:
        raise SystemExit(f"No hay PDFs en {data_dir}")
    return pdfs


def percentile(vals: list[float], p: float) -> float:
    if not vals:
        return float("nan")
    s = sorted(vals)
    if len(s) == 1:
        return s[0]
    k = (len(s) - 1) * (p / 100.0)
    f = int(k)
    c = min(f + 1, len(s) - 1)
    if f == c:
        return s[f]
    return s[f] + (s[c] - s[f]) * (k - f)


def count_one(base_url: str, pdf: Path, timeout_s: float) -> dict:
    t0 = time.perf_counter()
    status = "error"
    http_code = 0
    service_count = ""
    err = ""
    try:
        with pdf.open("rb") as fh:
            r = requests.post(
                f"{base_url}/count/v1/count",
                files={"files": (pdf.name, fh, "application/pdf")},
                timeout=timeout_s,
            )
        http_code = r.status_code
        try:
            body = r.json()
        except Exception:
            body = {}
        if r.status_code >= 500:
            status = f"http_{r.status_code}"
            err = (r.text or "")[:200].replace("\n", " ")
        elif r.status_code >= 400:
            status = f"http_{r.status_code}"
            err = str(body.get("error") or r.text)[:200]
        else:
            results = body.get("results") or []
            if results:
                service_count = results[0].get("service_count", "")
                status = str(results[0].get("status") or "ok")
            else:
                status = "ok_empty"
    except requests.Timeout:
        status = "timeout"
    except Exception as ex:
        status = "exception"
        err = str(ex)[:200]

    return {
        "file": pdf.name,
        "wall_s": round(time.perf_counter() - t0, 3),
        "http_code": http_code,
        "status": status,
        "service_count": service_count,
        "error": err,
    }


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="ST-COUNT-01 count-api load bench")
    p.add_argument("--base-url", default=os.environ.get("BASE_URL", "http://127.0.0.1:8189"))
    p.add_argument("--data-dir", type=Path, default=Path(os.environ.get("COUNT_BENCH_PDF_DIR") or DEFAULT_DATA))
    p.add_argument(
        "--profile",
        choices=sorted(PROFILES.keys()),
        default=os.environ.get("PROFILE") or "shared8g",
    )
    p.add_argument("--duration", type=float, default=float(os.environ.get("COUNT_BENCH_DURATION_S") or "0"))
    p.add_argument("--rate", type=float, default=float(os.environ.get("COUNT_BENCH_RATE") or "0"))
    p.add_argument("--max-workers", type=int, default=int(os.environ.get("COUNT_BENCH_MAX_WORKERS") or "0"))
    p.add_argument("--timeout", type=float, default=float(os.environ.get("COUNT_BENCH_TIMEOUT_S") or "120"))
    p.add_argument(
        "--out-csv",
        type=Path,
        default=Path(
            os.environ.get("COUNT_BENCH_CSV")
            or (DEFAULT_INFORMES / f"st_count_01_{date.today().strftime('%Y%m%d')}.csv")
        ),
    )
    return p.parse_args()


def main() -> int:
    args = parse_args()
    base_url = str(args.base_url).rstrip("/")
    cfg = PROFILES[args.profile]
    rate = args.rate if args.rate > 0 else cfg["rate"]
    duration_s = args.duration if args.duration > 0 else cfg["duration_s"]
    max_workers = args.max_workers if args.max_workers > 0 else cfg["max_workers"]
    interval = 1.0 / rate if rate > 0 else 1.0

    try:
        h = requests.get(f"{base_url}/health/live", timeout=5)
        print(f"health/live -> {h.status_code}")
    except Exception as ex:
        print(f"AVISO: health/live ({ex})", file=sys.stderr)

    pdfs = list_pdfs(args.data_dir)
    print(f"BASE_URL={base_url}")
    print(f"PROFILE={args.profile} rate={rate}/s duration={duration_s}s max_workers={max_workers}")
    print(f"DATA_DIR={args.data_dir} ({len(pdfs)} PDFs)")
    print(f"OUT={args.out_csv}")
    print()

    rows: list[dict] = []
    lock = threading.Lock()
    idx = {"i": 0}
    stop_at = time.perf_counter() + duration_s
    next_slot = time.perf_counter()

    def submit_one(ex: ThreadPoolExecutor, futures: set):
        nonlocal next_slot
        now = time.perf_counter()
        if now < next_slot:
            time.sleep(next_slot - now)
        next_slot = max(next_slot + interval, time.perf_counter())
        with lock:
            pdf = pdfs[idx["i"] % len(pdfs)]
            idx["i"] += 1
            n = idx["i"]
        print(f"[launch {n}] {pdf.name} ...", flush=True)
        futures.add(ex.submit(count_one, base_url, pdf, args.timeout))

    args.out_csv.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = ["file", "wall_s", "http_code", "status", "service_count", "error"]
    t0 = time.perf_counter()
    futures: set = set()

    with ThreadPoolExecutor(max_workers=max_workers) as ex:
        while time.perf_counter() < stop_at:
            # No saturar más allá de max_workers en vuelo
            while len(futures) >= max_workers:
                done = {f for f in futures if f.done()}
                for f in done:
                    row = f.result()
                    rows.append(row)
                    print(
                        f"  done wall_s={row['wall_s']} code={row['http_code']} status={row['status']}",
                        flush=True,
                    )
                    futures.remove(f)
                if len(futures) >= max_workers:
                    time.sleep(0.05)
            if time.perf_counter() >= stop_at:
                break
            submit_one(ex, futures)

        for f in as_completed(futures):
            row = f.result()
            rows.append(row)
            print(
                f"  done wall_s={row['wall_s']} code={row['http_code']} status={row['status']}",
                flush=True,
            )

    wall_total = round(time.perf_counter() - t0, 3)

    with args.out_csv.open("w", encoding="utf-8", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=fieldnames)
        w.writeheader()
        for row in rows:
            w.writerow(row)

    ok = [r for r in rows if r["http_code"] and 200 <= int(r["http_code"]) < 300]
    fail = [r for r in rows if r not in ok]
    times = [float(r["wall_s"]) for r in ok]
    meta = {
        "scenario": "ST-COUNT-01",
        "profile": args.profile,
        "rate_per_s": rate,
        "duration_s_planned": duration_s,
        "max_workers": max_workers,
        "n_total": len(rows),
        "n_ok": len(ok),
        "n_fail": len(fail),
        "wall_total_s": wall_total,
        "achieved_rate_per_s": round(len(rows) / wall_total, 3) if wall_total else 0,
        "csv": str(args.out_csv),
        "finished_at": datetime.now(timezone.utc).isoformat(),
    }
    if times:
        meta["wall_s_avg"] = round(statistics.mean(times), 3)
        meta["wall_s_p50"] = round(percentile(times, 50), 3)
        meta["wall_s_p95"] = round(percentile(times, 95), 3)
        meta["wall_s_min"] = round(min(times), 3)
        meta["wall_s_max"] = round(max(times), 3)
        # Throughput efectivo si se procesara en serie con p50 (referencia conservadora)
        meta["throughput_serial_p50_per_h"] = round(3600 / meta["wall_s_p50"], 1) if meta["wall_s_p50"] else None
        meta["throughput_parallel_est_per_h"] = (
            round(meta["throughput_serial_p50_per_h"] * max_workers, 1)
            if meta.get("throughput_serial_p50_per_h")
            else None
        )

    meta_path = args.out_csv.with_suffix(".summary.json")
    meta_path.write_text(json.dumps(meta, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

    print()
    print("=== Resumen ST-COUNT-01 (count-api carga) ===")
    print(f"CSV:      {args.out_csv}")
    print(f"Summary:  {meta_path}")
    print(f"ok={len(ok)} fail={len(fail)} total={len(rows)} wall_total={wall_total}s")
    print(f"achieved_rate ~ {meta['achieved_rate_per_s']}/s (plan {rate}/s)")
    if times:
        print(
            f"wall_s avg={meta['wall_s_avg']:.3f} p50={meta['wall_s_p50']:.3f} "
            f"p95={meta['wall_s_p95']:.3f} min={meta['wall_s_min']:.3f} max={meta['wall_s_max']:.3f}"
        )
        print(
            f"Throughput ref. serie(p50) ~ {meta['throughput_serial_p50_per_h']} fact/h; "
            f"x{max_workers} workers ~ {meta['throughput_parallel_est_per_h']} fact/h (ideal)"
        )
    oom_hint = any("timeout" in r["status"] or str(r["http_code"]).startswith("5") for r in rows)
    if oom_hint:
        print("AVISO: timeouts/5xx — posible contencion RAM/CPU; revisa docker stats.", file=sys.stderr)
    # PASS si hay ok y tasa de fallo < 15% (umbral checks k6 rate>0.85)
    fail_rate = len(fail) / max(len(rows), 1)
    return 0 if len(ok) > 0 and fail_rate <= 0.15 else 1


if __name__ == "__main__":
    sys.exit(main())