#!/usr/bin/env python3
"""
Hosted Checkout page-open load test: CloudLayer (4056) -> MastercardSimulator.

Opens N Hosted Checkout payment pages at the same moment, one per BULKTEST
merchant, exactly as a merchant server + payer browser would — and stops
before paying:

  1. POST /pgpayments/pay/init         X-Merchant-Id / X-Public-Key / X-Checksum
  2. GET  <payment_url>                CloudLayer runs INITIATE_CHECKOUT and
                                       returns the checkout page (session id)
  3. GET  simulator /static/checkout/session/<id>/context
                                       what checkout.min.js fetches in the browser

Two subcommands:

  prepare  Reads the bulk merchants' public/secret keys from MySQL (read-only)
           and writes a CSV of merchant_ref, public_key and the checksum for
           one fixed amount. The secret keys never leave memory.

             MYSQL_PWD=... python3 scripts/hosted_checkout_load_test.py prepare

  run      Opens pages in stages, e.g. 10, 50, 100 at once.

             python3 scripts/hosted_checkout_load_test.py run --stages 10,50,100

Every page opened creates a PENDING pg_transactions row for that bulk merchant,
which CloudLayer later marks EXPIRED (nobody pays). Stdlib only.
"""

import argparse
import csv
import hashlib
import hmac
import json
import os
import re
import statistics
import subprocess
import sys
import threading
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime

DEFAULT_CSV = "/home/priyanka/bulk_checkout_merchants.csv"
DEFAULT_BASE = "http://localhost:4056"
DEFAULT_SIM = "http://localhost:4444"
SESSION_RE = re.compile(r"session:\s*\{\s*id:\s*'([^']+)'")

MERCHANT_QUERY = """
SELECT c.user_id, c.merchant_id, m.merchant_refrence_number, u.public_key, u.secret_key
  FROM shukria_transactions.mpgs_merchant_credentials c
  JOIN shukria_mms_demo_local.users u         ON u.id = c.user_id
  JOIN shukria_mms_demo_local.user_metadata m ON m.user_id = c.user_id
 WHERE c.merchant_id LIKE 'BULKTEST%'
   AND c.method = 'hosted-checkout'
   AND c.active = 1
 ORDER BY c.merchant_id
"""


# ── prepare ──────────────────────────────────────────────────────────────────

def checksum(secret_key, merchant_ref, public_key, amount):
    # Same as CloudLayer's MerchantChecksum: hex(HMAC-SHA256(secret, "ref|public_key|amount"))
    source = f"{merchant_ref}|{public_key}|{amount}"
    return hmac.new(secret_key.encode(), source.encode(), hashlib.sha256).hexdigest()


def prepare(args):
    if "MYSQL_PWD" not in os.environ:
        sys.exit("Set MYSQL_PWD (the MySQL password from CloudLayer's config/dev.exs) first.")
    if os.path.exists(args.out):
        sys.exit(f"{args.out} already exists — delete it first if you want to regenerate it.")

    amount = f"{float(args.amount):.2f}"
    result = subprocess.run(
        ["mysql", "-u", args.db_user, "-h", args.db_host, "-N", "-B", "-e", MERCHANT_QUERY],
        capture_output=True, text=True, check=False,
    )
    if result.returncode != 0:
        sys.exit(f"mysql failed: {result.stderr.strip()}")

    rows = []
    for line in result.stdout.splitlines():
        user_id, merchant_id, ref, public_key, secret_key = line.split("\t")
        rows.append({
            "user_id": user_id,
            "merchant_id": merchant_id,
            "merchant_ref": ref,
            "public_key": public_key,
            "amount": amount,
            "checksum": checksum(secret_key, ref, public_key, amount),
        })

    if not rows:
        sys.exit("No active hosted-checkout BULKTEST merchants found.")

    fd = os.open(args.out, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
    with os.fdopen(fd, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        w.writeheader()
        w.writerows(rows)

    print(f"Wrote {len(rows)} merchants ({rows[0]['merchant_id']}..{rows[-1]['merchant_id']}) "
          f"for amount {amount} to {args.out} (mode 600)")


# ── run ──────────────────────────────────────────────────────────────────────

def http(method, url, body=None, headers=None, timeout=60):
    data = json.dumps(body).encode() if body is not None else None
    req = urllib.request.Request(url, data=data, method=method)
    for k, v in (headers or {}).items():
        req.add_header(k, v)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            return r.status, r.read()
    except urllib.error.HTTPError as e:
        return e.code, e.read()


def open_page(args, merchant, run_id, stage, barrier):
    order_id = f"HC{run_id}-S{stage}-{merchant['merchant_id']}"
    res = {"stage": stage, "merchant_id": merchant["merchant_id"], "order_id": order_id,
           "ok": False, "failed_step": "", "http": "", "error": "", "session_id": "",
           "init_ms": 0.0, "page_ms": 0.0, "context_ms": 0.0, "total_ms": 0.0}

    barrier.wait()  # every page in the stage opens together
    t0 = time.perf_counter()
    step = "init"
    try:
        # 1. merchant server -> CloudLayer: validate + get payment_url
        t = time.perf_counter()
        status, raw = http("POST", args.base + "/pgpayments/pay/init", {
            "amount": merchant["amount"],
            "currency": args.currency,
            "order_id": order_id,
            "customer_name": "Load Test",
            "email": "loadtest@example.com",
        }, {
            "Content-Type": "application/json",
            "X-Merchant-Id": merchant["merchant_ref"],
            "X-Public-Key": merchant["public_key"],
            "X-Checksum": merchant["checksum"],
        }, args.timeout)
        res["init_ms"] = (time.perf_counter() - t) * 1000
        data = json.loads(raw or b"{}") if status == 200 else {}
        if status != 200 or data.get("flow") != "mpgs" or not data.get("payment_url"):
            res.update(http=status, error=(data.get("flow") or raw[:100].decode(errors="replace")))
            raise RuntimeError

        # 2. payer browser -> CloudLayer: INITIATE_CHECKOUT + checkout page
        step = "page"
        t = time.perf_counter()
        status, html = http("GET", args.base + data["payment_url"], timeout=args.timeout)
        res["page_ms"] = (time.perf_counter() - t) * 1000
        match = SESSION_RE.search(html.decode(errors="replace")) if status == 200 else None
        if not match:
            res.update(http=status, error="no session id in page" if status == 200 else "")
            raise RuntimeError
        res["session_id"] = match.group(1)

        # 3. payer browser -> simulator: what checkout.min.js fetches
        step = "context"
        t = time.perf_counter()
        status, raw = http("GET", f"{args.sim}/static/checkout/session/{res['session_id']}/context",
                           timeout=args.timeout)
        res["context_ms"] = (time.perf_counter() - t) * 1000
        ctx = json.loads(raw or b"{}") if status == 200 else {}
        if status != 200 or not ctx.get("successIndicator"):
            res.update(http=status, error="no checkout context at simulator")
            raise RuntimeError

        res["ok"] = True
    except RuntimeError:
        res["failed_step"] = step
    except Exception as e:  # timeout, connection refused ...
        res.update(failed_step=step, error=f"{type(e).__name__}: {e}"[:120])

    res["total_ms"] = (time.perf_counter() - t0) * 1000
    for k in ("init_ms", "page_ms", "context_ms", "total_ms"):
        res[k] = round(res[k], 1)
    return res


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 summarise(stage, results, wall):
    ok = [r for r in results if r["ok"]]
    failed = [r for r in results if not r["ok"]]

    print(f"\n── Stage {stage}: {len(results)} pages opened at once ──")
    print(f"  pages ok     {len(ok)}  |  failed {len(failed)}")
    for label, key in (("total", "total_ms"), ("init", "init_ms"), ("page", "page_ms"), ("context", "context_ms")):
        vals = [r[key] for r in ok]
        if vals:
            print(f"  {label:<8} ms  p50 {pct(vals, 50):.0f}  p95 {pct(vals, 95):.0f}  "
                  f"max {max(vals):.0f}  mean {statistics.mean(vals):.0f}")
    print(f"  wall time    {wall:.1f}s  ->  {len(results) / wall:.1f} pages/s")

    if failed:
        groups = {}
        for r in failed:
            key = f"step={r['failed_step']} http={r['http'] or '-'} {r['error'] or ''}".strip()
            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])}")
    return len(failed)


def run(args):
    stages = [int(s) for s in args.stages.split(",") if s.strip()]
    with open(args.merchants, newline="") as f:
        merchants = sorted(csv.DictReader(f), key=lambda r: r["merchant_id"])

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

    print(f"CloudLayer : {args.base}/pgpayments/pay/init -> /pay/start")
    print(f"Simulator  : {args.sim}")
    print(f"Merchants  : {len(merchants)} in {args.merchants} (amount {merchants[0]['amount']} {args.currency})")
    print(f"Stages     : {stages}  ({sum(stages)} pages in total, {args.pause:.0f}s pause between stages)")
    print("Every page creates a PENDING pg_transactions row for a BULKTEST merchant (later EXPIRED).")

    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"checkout_results_{run_id}.csv"
    all_results = []

    for i, stage in enumerate(stages):
        barrier = threading.Barrier(stage)
        t0 = time.perf_counter()
        with ThreadPoolExecutor(max_workers=stage) as pool:
            results = list(pool.map(lambda m: open_page(args, m, run_id, stage, barrier), merchants[:stage]))
        wall = time.perf_counter() - t0
        all_results.extend(results)

        failures = summarise(stage, results, wall)

        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} failed pages (--stop-on-fail).")
            break
        if i < len(stages) - 1:
            time.sleep(args.pause)

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


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = p.add_subparsers(dest="cmd", required=True)

    pp = sub.add_parser("prepare", help="export merchant refs, public keys and checksums (needs MYSQL_PWD)")
    pp.add_argument("--out", default=DEFAULT_CSV)
    pp.add_argument("--amount", default="10.00", help="amount the checksums are computed for")
    pp.add_argument("--db-user", default="root")
    pp.add_argument("--db-host", default="localhost")

    pr = sub.add_parser("run", help="open checkout pages in stages")
    pr.add_argument("--merchants", default=DEFAULT_CSV)
    pr.add_argument("--base", default=DEFAULT_BASE, help="CloudLayer 4056 base URL")
    pr.add_argument("--sim", default=DEFAULT_SIM, help="simulator base URL")
    pr.add_argument("--stages", default="10,50,100")
    pr.add_argument("--pause", type=float, default=15)
    pr.add_argument("--timeout", type=float, default=60)
    pr.add_argument("--currency", default="AED")
    pr.add_argument("--stop-on-fail", action="store_true")
    pr.add_argument("--yes", action="store_true")
    pr.add_argument("--out", default=None, help="results CSV (default: checkout_results_<timestamp>.csv)")

    args = p.parse_args()
    prepare(args) if args.cmd == "prepare" else run(args)


if __name__ == "__main__":
    main()