#!/usr/bin/env python3
"""
Staged multi-merchant load test: CloudLayer Direct API -> MastercardSimulator.

Each stage sends ONE payment per merchant for the first N merchants in the
keys CSV, all fired at the same moment (a barrier releases every thread
together), then checks the simulator's /admin/stats to confirm every
merchant's payment actually arrived there.

    python3 scripts/bulk_load_test.py --stages 10,50,100,250,500,1000

Stdlib only. API keys are read from the CSV and never printed or written to
the results file.
"""

import argparse
import base64
import csv
import json
import statistics
import sys
import threading
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime

DEFAULT_URL = "https://mercurypay.aritic.app/pn/pgpayments/api/direct/pay"
DEFAULT_SIM = "http://localhost:4444"
SIM_USER = "merchant.BULKTEST0001"
SIM_PASS = "bulk_password_123"


def parse_args():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--keys", default="/home/priyanka/bulk_keys.csv", help="bulk_keys.csv (user_id,email,merchant_id,api_key)")
    p.add_argument("--url", default=DEFAULT_URL, help="CloudLayer Direct /pay endpoint (4056 instance)")
    p.add_argument("--stages", default="10,50,100,250,500,1000", help="comma-separated merchant counts")
    p.add_argument("--pause", type=float, default=10, help="seconds to wait between stages")
    p.add_argument("--timeout", type=float, default=60, help="per-request timeout in seconds")
    p.add_argument("--card", default="5123450000000008", help="test card (default approves)")
    p.add_argument("--amount", default="10.00")
    p.add_argument("--currency", default="AED")
    p.add_argument("--sim", default=DEFAULT_SIM, help="simulator base URL, for /admin/stats checks")
    p.add_argument("--stop-on-fail", action="store_true", help="stop before the next stage if this one had any failure")
    p.add_argument("--yes", action="store_true", help="don't ask for confirmation")
    p.add_argument("--out", default=None, help="results CSV (default: load_results_<timestamp>.csv)")
    return p.parse_args()


def load_merchants(path):
    with open(path, newline="") as f:
        rows = [r for r in csv.DictReader(f) if r.get("api_key") and r.get("merchant_id")]
    rows.sort(key=lambda r: r["merchant_id"])
    return rows


def sim_get(sim, path):
    req = urllib.request.Request(sim + path)
    token = base64.b64encode(f"{SIM_USER}:{SIM_PASS}".encode()).decode()
    req.add_header("Authorization", "Basic " + token)
    with urllib.request.urlopen(req, timeout=30) as r:
        return json.load(r)


def pay(args, merchant, run_id, stage, barrier):
    order_id = f"LT{run_id}-S{stage}-{merchant['merchant_id']}"
    body = json.dumps({
        "amount": args.amount,
        "currency": args.currency,
        "order_id": order_id,
        "customer_name": "Load Test",
        "card": {
            "number": args.card,
            "expiry_month": "12",
            "expiry_year": "27",
            "security_code": "123",
            "name_on_card": "Load Test",
        },
    }).encode()

    req = urllib.request.Request(args.url, data=body, method="POST")
    req.add_header("Content-Type", "application/json")
    req.add_header("Authorization", "Bearer " + merchant["api_key"])

    result = {"stage": stage, "merchant_id": merchant["merchant_id"], "order_id": order_id,
              "http": 0, "status": "", "status_code": "", "message": "", "latency_ms": 0.0}

    barrier.wait()  # every thread in the stage fires together
    start = time.perf_counter()
    try:
        with urllib.request.urlopen(req, timeout=args.timeout) as r:
            result["http"] = r.status
            data = json.loads(r.read() or b"{}")
    except urllib.error.HTTPError as e:
        result["http"] = e.code
        try:
            data = json.loads(e.read() or b"{}")
        except ValueError:
            data = {}
    except Exception as e:  # timeout, connection refused, TLS ...
        data = {"message": f"{type(e).__name__}: {e}"}
    result["latency_ms"] = round((time.perf_counter() - start) * 1000, 1)

    result["status"] = str(data.get("status", ""))
    result["status_code"] = str(data.get("status_code", ""))
    result["message"] = str(data.get("message", data.get("error", "")))[:120]
    return result


def pct(values, q):
    if not values:
        return 0.0
    values = sorted(values)
    return values[min(len(values) - 1, int(round(q / 100 * (len(values) - 1))))]


def run_stage(args, merchants, run_id, stage):
    barrier = threading.Barrier(len(merchants))
    t0 = time.perf_counter()
    with ThreadPoolExecutor(max_workers=len(merchants)) as pool:
        results = list(pool.map(lambda m: pay(args, m, run_id, stage, barrier), merchants))
    wall = time.perf_counter() - t0
    return results, wall


def summarise(stage, results, wall, before, after, sim_ok):
    ok = [r for r in results if r["status"] == "success"]
    failed = [r for r in results if r["status"] != "success"]
    lat = [r["latency_ms"] for r in results]

    print(f"\n── Stage {stage}: {len(results)} merchants ──")
    print(f"  CloudLayer   success {len(ok)}  |  not success {len(failed)}")
    print(f"  latency ms   p50 {pct(lat, 50):.0f}  p95 {pct(lat, 95):.0f}  max {max(lat):.0f}  "
          f"mean {statistics.mean(lat):.0f}")
    print(f"  wall time    {wall:.1f}s  ->  {len(results) / wall:.1f} payments/s")

    if failed:
        groups = {}
        for r in failed:
            key = f"http={r['http']} status={r['status'] or '-'} msg={r['message'] or '-'}"
            groups.setdefault(key, []).append(r["merchant_id"])
        print("  failures:")
        for key, mids in sorted(groups.items(), key=lambda kv: -len(kv[1])):
            print(f"    {len(mids):>4} × {key}   e.g. {', '.join(mids[:3])}")

    if sim_ok:
        delta = after["total"] - before["total"]
        missing = [r["merchant_id"] for r in results
                   if sum(after["by_merchant"].get(r["merchant_id"], {}).values())
                   - sum(before["by_merchant"].get(r["merchant_id"], {}).values()) < 1]
        status_delta = {k: v - before["by_status"].get(k, 0)
                        for k, v in after["by_status"].items() if v - before["by_status"].get(k, 0)}
        print(f"  simulator    +{delta} transactions {status_delta}")
        if missing:
            print(f"  ⚠ {len(missing)} merchants have no new simulator transaction, e.g. {', '.join(missing[:5])}")
        else:
            print("  ✓ every merchant in this stage reached the simulator")
    return len(failed)


def main():
    args = parse_args()
    stages = [int(s) for s in args.stages.split(",") if s.strip()]
    merchants = load_merchants(args.keys)

    if max(stages) > len(merchants):
        sys.exit(f"Largest stage is {max(stages)} but the CSV only has {len(merchants)} merchants")

    total = sum(stages)
    print(f"Endpoint : {args.url}")
    print(f"Merchants: {len(merchants)} in {args.keys}")
    print(f"Stages   : {stages}  ({total} payments in total, {args.pause:.0f}s pause between stages)")
    print(f"Card     : {args.card[:6]}…{args.card[-4:]}  amount {args.amount} {args.currency}")
    print("Every payment is a real CloudLayer transaction (pg_transactions) for a BULKTEST merchant.")

    if not args.yes and input("Proceed? [y/N] ").strip().lower() != "y":
        sys.exit("Aborted.")

    run_id = datetime.now().strftime("%m%d%H%M%S")
    out = args.out or f"load_results_{run_id}.csv"
    all_results = []

    for i, stage in enumerate(stages):
        try:
            before = sim_get(args.sim, "/admin/stats")
            sim_ok = True
        except Exception as e:
            print(f"  (simulator stats unavailable: {e})")
            before, sim_ok = None, False

        results, wall = run_stage(args, merchants[:stage], run_id, stage)
        all_results.extend(results)

        after = None
        if sim_ok:
            try:
                after = sim_get(args.sim, "/admin/stats")
            except Exception as e:
                print(f"  (simulator stats unavailable: {e})")
                sim_ok = False

        failures = summarise(stage, results, wall, before, after, sim_ok)

        with open(out, "w", newline="") as f:
            w = csv.DictWriter(f, fieldnames=list(all_results[0].keys()))
            w.writeheader()
            w.writerows(all_results)

        if failures and args.stop_on_fail:
            print(f"\nStopping: stage {stage} had {failures} failures (--stop-on-fail).")
            break
        if i < len(stages) - 1:
            time.sleep(args.pause)

    print(f"\nPer-request results: {out}")


if __name__ == "__main__":
    main()