#!/usr/bin/env python3
"""Compare captured mitmproxy requests between claude-code and claude-py."""

import json
import sys
from pathlib import Path

CAPTURES_DIR = sys.argv[1] if len(sys.argv) > 1 else "./test/captures"

IGNORE_HEADERS = {
    "authorization",
    "content-length",
    "host",
    "connection",
    "proxy-connection",
    "proxy-authorization",
    "transfer-encoding",
    "accept-encoding",
}

VOLATILE_BODY_KEYS = {"messages"}
VOLATILE_METADATA_KEYS = {"user_id"}


def load_captures(directory):
    caps = []
    for f in sorted(Path(directory).glob("req_*.json")):
        with open(f) as fh:
            caps.append(json.load(fh))
    return caps


def normalize_headers(headers):
    return {k.lower(): v for k, v in headers.items() if k.lower() not in IGNORE_HEADERS}


def normalize_body(body):
    if not isinstance(body, dict):
        return body
    out = {}
    for k, v in body.items():
        if k in VOLATILE_BODY_KEYS:
            out[k] = f"<{len(v)} items>" if isinstance(v, list) else "<redacted>"
        elif k == "metadata":
            out[k] = {
                mk: mv
                for mk, mv in (v if isinstance(v, dict) else {}).items()
                if mk not in VOLATILE_METADATA_KEYS
            }
        else:
            out[k] = v
    return out


def fmt(v):
    if isinstance(v, (dict, list)):
        return json.dumps(v, indent=2, default=str)
    return str(v)


def compare_dicts(label, d1, d2, n1, n2):
    keys = sorted(set(list(d1.keys()) + list(d2.keys())))
    diffs = 0
    for key in keys:
        v1 = d1.get(key, "<missing>")
        v2 = d2.get(key, "<missing>")
        if v1 == v2:
            print(f"    \033[32m✓\033[0m {key}")
        else:
            diffs += 1
            print(f"    \033[31m✗\033[0m {key}")
            print(f"        {n1}: {fmt(v1)}")
            print(f"        {n2}: {fmt(v2)}")
    if diffs == 0:
        print(f"  \033[32m{label}: IDENTICAL\033[0m")
    else:
        print(f"  \033[31m{label}: {diffs} difference(s)\033[0m")
    return diffs


def dump_request(cap):
    src = cap.get("source", "?")
    print(f"  Source: {src}")
    print(f"  URL:    {cap['url']}")
    headers = normalize_headers(cap.get("headers", {}))
    body = cap.get("body", {})
    print(f"  Headers ({len(headers)}):")
    for k in sorted(headers):
        print(f"    {k}: {headers[k]}")
    if isinstance(body, dict):
        print(f"  Body keys: {sorted(body.keys())}")
        print(f"    model:      {body.get('model', 'N/A')}")
        print(f"    stream:     {body.get('stream', 'N/A')}")
        print(f"    max_tokens: {body.get('max_tokens', 'N/A')}")
        tools = body.get("tools", [])
        if isinstance(tools, list):
            print(f"    tools ({len(tools)}): {[t.get('name', '?') for t in tools]}")
        system = body.get("system", [])
        if isinstance(system, list):
            print(f"    system blocks: {len(system)}")


def main():
    caps = load_captures(CAPTURES_DIR)
    if not caps:
        print(f"No captures found in {CAPTURES_DIR}")
        print("Run claude-code and claude-py through the proxy first.")
        sys.exit(1)

    by_source = {}
    for c in caps:
        src = c.get("source", "unknown")
        by_source.setdefault(src, []).append(c)

    print(f"Found {len(caps)} captured request(s)")
    print(f"Sources: { {k: len(v) for k, v in by_source.items()} }\n")

    for i, cap in enumerate(caps):
        print(f"--- Request #{cap.get('capture_id', i + 1)} ---")
        dump_request(cap)
        print()

    sources = list(by_source.keys())
    if len(sources) == 2 and all(len(v) >= 1 for v in by_source.values()):
        n1, n2 = sources
        c1, c2 = by_source[n1][0], by_source[n2][0]

        print("=" * 60)
        print(f"COMPARING: {n1} vs {n2}")
        print("=" * 60)

        h1 = normalize_headers(c1.get("headers", {}))
        h2 = normalize_headers(c2.get("headers", {}))
        total = compare_dicts("Headers", h1, h2, n1, n2)

        b1 = normalize_body(c1.get("body", {}))
        b2 = normalize_body(c2.get("body", {}))
        total += compare_dicts("Body", b1, b2, n1, n2)

        t1 = c1.get("body", {}).get("tools", [])
        t2 = c2.get("body", {}).get("tools", [])
        if isinstance(t1, list) and isinstance(t2, list):
            tn1 = sorted(t.get("name", "") for t in t1)
            tn2 = sorted(t.get("name", "") for t in t2)
            if tn1 == tn2:
                print(f"  \033[32mTool names: IDENTICAL ({len(tn1)} tools)\033[0m")
            else:
                total += 1
                only1 = set(tn1) - set(tn2)
                only2 = set(tn2) - set(tn1)
                print("  \033[31mTool names: DIFFER\033[0m")
                if only1:
                    print(f"    Only in {n1}: {only1}")
                if only2:
                    print(f"    Only in {n2}: {only2}")

        print(f"\n{'=' * 60}")
        if total == 0:
            print("\033[32mRESULT: Requests are identical on the wire.\033[0m")
        else:
            print(f"\033[31mRESULT: {total} total difference(s).\033[0m")
        sys.exit(0 if total == 0 else 1)
    else:
        print("Tip: Clear test/captures/ and run one request from each client to compare.")


if __name__ == "__main__":
    main()