#!/usr/bin/env python
"""Build the frozen request manifest from the recorder's raw output.

Reads every ``raw_*.jsonl`` in ``--out`` (one per xdist worker), dedupes the
recorded calls into a stable, ordered set of unique requests, and writes:

* ``manifest.json``             -- the frozen request set; each request gets a
  stable zero-padded ``id``. This is the reusable artifact: ``replay.py`` fires
  exactly this set against any server build.
* ``capture_after_extract.json`` -- the responses observed *during extraction*
  (server at the after-ref), in the same schema ``replay.py`` emits. Because a
  request
  may be exercised by several xdist workers, a request seen with two different
  responses is flagged ``intra_run_distinct_sha`` -- a free nondeterminism
  signal independent of the before/after comparison.
"""
import argparse
import base64
import datetime
import glob
import hashlib
import json
import os
import subprocess
from collections import Counter
from urllib.parse import urlsplit

from requests.models import PreparedRequest

HERE = os.path.dirname(os.path.abspath(__file__))
# Defensive: only dev-server hosts belong in the manifest (the recorder already
# filters, but stale raw files or a recorder bug should not leak other traffic).
TARGET_HOSTS = {
    h.strip()
    for h in os.environ.get("ENDPOINT_DIFF_TARGET_HOSTS", "localhost,127.0.0.1").split(
        ","
    )
    if h.strip()
}


def canonical_url(url, params):
    """The exact URL ``requests`` would send (params folded into the query)."""
    if not params:
        return url
    pr = PreparedRequest()
    pr.prepare_url(url, params)
    return pr.url


def request_key(method, url, headers, json_body, data):
    """Stable identity of a request: method + URL + headers + body."""
    blob = json.dumps(
        {
            "method": method,
            "url": url,
            "headers": headers,
            "json": json_body,
            "data": data,
        },
        sort_keys=True,
        default=str,
    )
    return hashlib.sha256(blob.encode()).hexdigest()


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

    Kept byte-identical to ``deep_sort`` in replay.py -- the two producers must
    agree or their ``sha256_sorted`` hashes are not comparable across captures.
    """
    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 main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default=os.path.join(HERE, "_out"))
    args = ap.parse_args()

    raw_files = sorted(glob.glob(os.path.join(args.out, "raw_*.jsonl")))
    if not raw_files:
        raise SystemExit(f"no raw_*.jsonl found in {args.out} -- run extraction first")

    entries = {}  # key -> manifest entry (method/url/headers/json/data)
    observations = {}  # key -> list of observed responses
    total = skipped = 0

    for path in raw_files:
        with open(path) as fh:
            for line in fh:
                line = line.strip()
                if not line:
                    continue
                rec = json.loads(line)
                if rec.get("record_error"):
                    skipped += 1
                    continue
                if urlsplit(str(rec["url"])).hostname not in TARGET_HOSTS:
                    skipped += 1
                    continue
                total += 1
                method = rec["method"]
                url = canonical_url(rec["url"], rec.get("params"))
                headers = rec.get("headers") or {}
                json_body = rec.get("json")
                data = rec.get("data")
                key = request_key(method, url, headers, json_body, data)
                if key not in entries:
                    entries[key] = {
                        "method": method,
                        "url": url,
                        "headers": headers,
                        "json": json_body,
                        "data": data,
                    }
                    observations[key] = []
                if rec.get("resp_body_b64") is not None:
                    raw = base64.b64decode(rec["resp_body_b64"])
                    text = None
                else:
                    text = rec.get("resp_body") or ""
                    raw = text.encode("utf-8")
                canon, ordered = canonical_hashes(text)
                observations[key].append(
                    {
                        "status": rec.get("status"),
                        "content_type": rec.get("content_type"),
                        "len": rec.get("resp_len"),
                        "sha256": hashlib.sha256(raw).hexdigest(),
                        "sha256_canon": canon,
                        "sha256_sorted": ordered,
                    }
                )

    ordered = sorted(
        entries.items(),
        key=lambda kv: (
            kv[1]["method"],
            kv[1]["url"],
            json.dumps(kv[1]["headers"], sort_keys=True),
            json.dumps(kv[1]["json"], sort_keys=True, default=str),
        ),
    )

    manifest_requests = []
    extract_responses = {}
    for idx, (key, entry) in enumerate(ordered):
        rid = f"{idx:05d}"
        manifest_requests.append({"id": rid, **entry})
        obs = observations[key]
        distinct = sorted({o["sha256"] for o in obs})
        first = obs[0]
        extract_responses[rid] = {
            "status": first["status"],
            "content_type": first["content_type"],
            "len": first["len"],
            "sha256": first["sha256"],
            "sha256_canon": first["sha256_canon"],
            "sha256_sorted": first["sha256_sorted"],
            "body": None,
            "body_b64": None,
            "elapsed_ms": None,
            "error": None,
            "observations": len(obs),
            "intra_run_distinct_sha": distinct if len(distinct) > 1 else None,
        }

    try:
        git_rev = subprocess.check_output(
            ["git", "rev-parse", "HEAD"], text=True, cwd=HERE
        ).strip()
    except Exception:
        git_rev = "?"
    now = datetime.datetime.now(datetime.timezone.utc).isoformat()

    with open(os.path.join(args.out, "manifest.json"), "w") as fh:
        json.dump(
            {
                "generated_at": now,
                "git_rev": git_rev,
                "total_recorded": total,
                "unique_requests": len(manifest_requests),
                "requests": manifest_requests,
            },
            fh,
            indent=2,
            sort_keys=True,
        )

    with open(os.path.join(args.out, "capture_after_extract.json"), "w") as fh:
        json.dump(
            {
                "meta": {
                    "label": "after_extract",
                    "mechanism": "extract",
                    "git_rev": git_rev,
                    "finished_at": now,
                    "request_count": len(extract_responses),
                    "bodies": False,
                },
                "responses": extract_responses,
            },
            fh,
            indent=2,
            sort_keys=True,
        )

    by_method = Counter(r["method"] for r in manifest_requests)
    nondet = sum(1 for r in extract_responses.values() if r["intra_run_distinct_sha"])
    print(f"raw files        : {len(raw_files)}")
    print(f"records read     : {total}  (skipped {skipped} with record_error)")
    print(f"unique requests  : {len(manifest_requests)}  {dict(by_method)}")
    print(f"intra-extraction nondeterministic responses: {nondet}")
    print(f"wrote {args.out}/manifest.json")
    print(f"wrote {args.out}/capture_after_extract.json")


if __name__ == "__main__":
    main()