#!/usr/bin/env python
"""Replay a frozen endpoint_diff manifest against a server, capture every
response byte-for-byte.

Standalone on purpose: the only third-party import is ``requests`` (always in
the project venv). It does **not** import ``analytics`` or any test module, so
it runs unchanged whatever git revision is checked out -- which is what lets the
same request set be fired at the "before" and "after" builds.

    replay.py --label before --base-url http://127.0.0.1:5000

Writes ``capture_<label>.json``: ``meta`` + ``responses`` keyed by manifest id.
Each response carries status, Content-Type, byte length, round-trip time and
three sha256 hashes of the body -- ``sha256`` (raw bytes), ``sha256_canon``
(key-sorted JSON: equal iff the JSON value is equal) and ``sha256_sorted``
(key- and list-element-sorted JSON: equal iff the data matches up to array
order). The hashes alone let ``compare.py`` classify every difference, so the
body itself is stored only under ``--keep-bodies`` -- which makes captures far
larger and is meant for a small or targeted manifest, not the full suite.
"""
import argparse
import base64
import concurrent.futures
import datetime
import hashlib
import json
import os
import subprocess
import sys
import threading
from urllib.parse import urlsplit, urlunsplit

import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

HERE = os.path.dirname(os.path.abspath(__file__))
_local = threading.local()


def session():
    """One retrying Session per worker thread (Sessions are not thread-safe).

    Retry policy mirrors tests/integration/request.py so replay fidelity matches
    what the suite itself would have done.
    """
    s = getattr(_local, "session", None)
    if s is None:
        s = requests.Session()
        retry = Retry(total=2, backoff_factor=2, status_forcelist=[500, 502, 503, 504])
        adapter = HTTPAdapter(max_retries=retry)
        s.mount("http://", adapter)
        s.mount("https://", adapter)
        _local.session = s
    return s


def rebase(url, base_url):
    """Swap scheme+host of ``url`` for those of ``base_url``; keep path+query."""
    u, b = urlsplit(url), urlsplit(base_url)
    return urlunsplit((b.scheme, b.netloc, u.path, u.query, u.fragment))


def deep_sort(obj):
    """Canonical form with every dict key-sorted and every list element-sorted.

    Kept byte-identical to ``deep_sort`` in build_manifest.py -- the two
    producers must agree or their ``sha256_sorted`` hashes are not comparable.
    """
    if isinstance(obj, dict):
        return {k: deep_sort(obj[k]) for k in sorted(obj)}
    if isinstance(obj, list):
        return sorted(
            (deep_sort(x) for x in obj),
            key=lambda v: json.dumps(v, sort_keys=True, default=str),
        )
    return obj


def canonical_hashes(text):
    """``(sha256_canon, sha256_sorted)`` for a response body.

    ``canon``  -- sha256 of the key-sorted JSON re-serialization (list order
    kept): two responses with equal ``canon`` carry the same JSON value.
    ``sorted`` -- sha256 of the same with every list element-sorted too: equal
    ``sorted`` means the same data up to array element-order.
    Returns ``(None, None)`` when the body is absent or is not JSON.
    """
    if text is None:
        return None, None
    try:
        obj = json.loads(text)
    except (ValueError, TypeError):
        return None, None
    canon = hashlib.sha256(
        json.dumps(obj, sort_keys=True, separators=(",", ":"), default=str).encode()
    ).hexdigest()
    ordered = hashlib.sha256(
        json.dumps(
            deep_sort(obj), sort_keys=True, separators=(",", ":"), default=str
        ).encode()
    ).hexdigest()
    return canon, ordered


def fire(req, base_url, timeout, keep_bodies):
    url = rebase(req["url"], base_url)
    kwargs = {"timeout": timeout, "allow_redirects": False}
    if req.get("headers"):
        kwargs["headers"] = req["headers"]
    if req.get("json") is not None:
        kwargs["json"] = req["json"]
    if req.get("data") is not None:
        kwargs["data"] = req["data"]
    try:
        resp = session().request(req["method"], url, **kwargs)
    except Exception as exc:
        return req["id"], {
            "status": None,
            "content_type": None,
            "len": None,
            "sha256": None,
            "sha256_canon": None,
            "sha256_sorted": None,
            "body": None,
            "body_b64": None,
            "elapsed_ms": None,
            "error": repr(exc),
        }
    raw = resp.content
    try:
        text = raw.decode("utf-8")
    except UnicodeDecodeError:
        text = None
    canon, ordered = canonical_hashes(text)
    record = {
        "status": resp.status_code,
        "content_type": resp.headers.get("Content-Type"),
        "len": len(raw),
        "sha256": hashlib.sha256(raw).hexdigest(),
        "sha256_canon": canon,
        "sha256_sorted": ordered,
        "body": None,
        "body_b64": None,
        "elapsed_ms": round(resp.elapsed.total_seconds() * 1000, 1),
        "error": None,
    }
    if keep_bodies:
        if text is None:
            record["body_b64"] = base64.b64encode(raw).decode("ascii")
        else:
            record["body"] = text
    return req["id"], record


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument(
        "--label", required=True, help="capture label, e.g. before / after_replay"
    )
    ap.add_argument("--base-url", default="http://127.0.0.1:5000")
    ap.add_argument("--manifest", default=None)
    ap.add_argument("--out", default=None)
    ap.add_argument("--concurrency", type=int, default=15)
    ap.add_argument("--timeout", type=float, default=300.0)
    ap.add_argument("--git-rev", default=None)
    ap.add_argument(
        "--keep-bodies",
        action="store_true",
        help="also store full response bodies -- large captures; needed only "
        "for compare.py line-level body diffs",
    )
    args = ap.parse_args()

    out = args.out or os.path.join(HERE, "_out")
    manifest_path = args.manifest or os.path.join(out, "manifest.json")
    with open(manifest_path) as fh:
        manifest = json.load(fh)
    reqs = manifest["requests"]
    total = len(reqs)

    git_rev = args.git_rev
    if git_rev is None:
        try:
            git_rev = subprocess.check_output(
                ["git", "rev-parse", "HEAD"], text=True, cwd=HERE
            ).strip()
        except Exception:
            git_rev = "?"

    started = datetime.datetime.now(datetime.timezone.utc).isoformat()
    responses = {}
    done = 0
    lock = threading.Lock()
    print(
        f"[{args.label}] replaying {total} requests at {args.base_url} "
        f"(concurrency={args.concurrency})",
        file=sys.stderr,
        flush=True,
    )
    with concurrent.futures.ThreadPoolExecutor(max_workers=args.concurrency) as ex:
        futures = [
            ex.submit(fire, r, args.base_url, args.timeout, args.keep_bodies)
            for r in reqs
        ]
        for future in concurrent.futures.as_completed(futures):
            rid, result = future.result()
            with lock:
                responses[rid] = result
                done += 1
                if done % 100 == 0 or done == total:
                    print(
                        f"  [{args.label}] {done}/{total}", file=sys.stderr, flush=True
                    )
    finished = datetime.datetime.now(datetime.timezone.utc).isoformat()

    capture = {
        "meta": {
            "label": args.label,
            "mechanism": "replay",
            "base_url": args.base_url,
            "git_rev": git_rev,
            "manifest_git_rev": manifest.get("git_rev"),
            "started_at": started,
            "finished_at": finished,
            "request_count": total,
            "concurrency": args.concurrency,
            "bodies": args.keep_bodies,
        },
        "responses": responses,
    }
    path = os.path.join(out, f"capture_{args.label}.json")
    with open(path, "w") as fh:
        json.dump(capture, fh, indent=2, sort_keys=True)

    statuses = {}
    errors = 0
    for r in responses.values():
        statuses[r["status"]] = statuses.get(r["status"], 0) + 1
        if r["error"]:
            errors += 1
    print(f"[{args.label}] wrote {path}")
    print(
        f"[{args.label}] status breakdown: {dict(sorted(statuses.items(), key=str))}"
        f"  transport errors: {errors}"
    )


if __name__ == "__main__":
    main()