"""
Send a burst of requests to the FastAPI endpoint to populate the Grafana dashboard.

Usage:
    python scripts/load_test.py               # default: 30 requests
    python scripts/load_test.py --requests 60 # more requests
    python scripts/load_test.py --url http://localhost:8000  # custom URL

Requires: pip install requests
"""
import os
import sys
import time
import random
import argparse
import requests

SAMPLES_DIR = os.path.join(os.path.dirname(__file__), "..", "samples")
NOISE_TYPES = ["gaussian", "salt_pepper", "speckle"]


def get_images():
    images = [
        os.path.join(SAMPLES_DIR, f)
        for f in os.listdir(SAMPLES_DIR)
        if f.endswith(".png")
    ]
    if not images:
        print("No sample images found. Run: python scripts/generate_samples.py")
        sys.exit(1)
    return images


def send_request(url: str, image_path: str, noise_type: str) -> dict:
    with open(image_path, "rb") as f:
        files = {"file": (os.path.basename(image_path), f, "image/png")}
        data  = {"noise_type": noise_type}
        t0 = time.perf_counter()
        resp = requests.post(f"{url}/predict", files=files, data=data, timeout=60)
        elapsed = time.perf_counter() - t0

    return {
        "status": resp.status_code,
        "noise":  noise_type,
        "image":  os.path.basename(image_path),
        "time":   elapsed,
        "ok":     resp.status_code == 200,
    }


def run(url: str, n_requests: int):
    images = get_images()
    print(f"Sending {n_requests} requests to {url}")
    print(f"Images: {[os.path.basename(i) for i in images]}")
    print("-" * 55)

    ok = err = 0

    for i in range(1, n_requests + 1):
        image      = random.choice(images)
        noise_type = random.choice(NOISE_TYPES)

        try:
            r = send_request(url, image, noise_type)
            status = "OK " if r["ok"] else "ERR"
            print(f"[{i:3d}/{n_requests}] {status}  {noise_type:12s}  "
                  f"{r['image']:20s}  {r['time']:.2f}s")
            if r["ok"]:
                ok += 1
            else:
                err += 1
        except Exception as e:
            print(f"[{i:3d}/{n_requests}] FAIL  {noise_type:12s}  {e}")
            err += 1

        # Vary delay so Grafana graphs look natural, not like a flat line
        delay = random.uniform(0.3, 1.2)
        time.sleep(delay)

    print("-" * 55)
    print(f"Done — {ok} success / {err} errors out of {n_requests} requests")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--url",      default="http://localhost:8000")
    parser.add_argument("--requests", type=int, default=30)
    args = parser.parse_args()
    run(args.url, args.requests)