#!/usr/bin/env python3
"""Bench and canonical dump for the 2026-10-13 rectification scoring speedup.

Read-only aside from the dump files the caller asks for. Cross-process dumps
must be launched with PYTHONHASHSEED=0. The baseline tree has no
include_internal parameter; that call falls back to the old signature.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import statistics
import sys
import time
from datetime import date
from pathlib import Path
from typing import Any, Callable

ROOT = Path(__file__).resolve().parents[2]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from scripts.rectification.api_service import block_scan, diagnostics, score_candidates
from scripts.rectification.contracts import normalize_rectification_request
from scripts.rectification import event_probes
from scripts.rectification import scoring_service
from scripts import active_rectification_event_engine as event_engine

FROZEN_TODAY = date(2026, 9, 16)
EVENT_SPECS = (
    ("education", "education_start", "2008-01-01", "2008-12-31", "year"),
    ("career", "career_entry", "2012-06-01", "2012-06-30", "month"),
    ("relocation", "relocation", "2014-03-01", "2014-03-31", "month"),
    ("relationship", "relationship_commitment", "2016-10-01", "2016-10-31", "month"),
    ("career", "promotion", "2018-01-01", "2018-12-31", "year"),
    ("finance", "asset_change", "2019-05-01", "2019-05-31", "month"),
    ("family", "family_event", "2020-08-01", "2020-08-31", "month"),
    ("health", "self_health_event", "2022-01-01", "2022-12-31", "year"),
)


class _FrozenDate(date):
    @classmethod
    def today(cls) -> date:
        return FROZEN_TODAY


def _events() -> list[dict[str, Any]]:
    rows = []
    for index, (domain, kind, start, end, precision) in enumerate(EVENT_SPECS, start=1):
        rows.append({
            "id": f"00000000-0000-4000-8000-{index:012d}",
            "domain": domain,
            "event_kind": kind,
            "date_start": start,
            "date_end": end,
            "precision": precision,
            "summary": "虚构",
        })
    return rows


def make_request(start: str, end: str, **extra: Any) -> dict[str, Any]:
    body = {
        "birth_date": "1990-01-01",
        "start_time": start,
        "end_time": end,
        "lat": 39.9,
        "lon": 116.4,
        "tz": 8,
        "events": _events(),
        **extra,
    }
    return normalize_rectification_request(body, today=FROZEN_TODAY)


def _strip(value: Any, keys: set[str]) -> Any:
    if isinstance(value, dict):
        return {key: _strip(item, keys) for key, item in value.items() if key not in keys}
    if isinstance(value, list):
        return [_strip(item, keys) for item in value]
    return value


def _canonical(value: Any) -> bytes:
    return json.dumps(value, sort_keys=True, ensure_ascii=True, separators=(",", ":")).encode("ascii")


def _score(request: dict[str, Any]) -> dict[str, Any]:
    try:
        return score_candidates(request, include_internal=True)
    except TypeError:
        return score_candidates(request)


def _call(kind: str, request: dict[str, Any], *, internal: bool = False) -> dict[str, Any]:
    if kind == "score":
        # Bench times the product call. Dump asks for the internal matrix so
        # the baseline tree, which always returned it, stays comparable.
        if internal:
            return _score(request)
        return score_candidates(request)
    if kind == "block":
        return block_scan(request)
    if kind == "diagnostics":
        return diagnostics(request)
    raise ValueError(kind)


def scenarios() -> list[dict[str, Any]]:
    from tests.test_rectification_engine_memoization import public_score_request

    rows = [
        {"name": "12:00-12:20", "kind": "score", "request": make_request("12:00", "12:20")},
        {"name": "11:45-12:15", "kind": "score", "request": make_request("11:45", "12:15")},
        {"name": "11:00-13:00", "kind": "score", "request": make_request("11:00", "13:00")},
        {"name": "23:20-00:40", "kind": "score", "request": make_request("23:20", "00:40")},
        {"name": "declined-relationship-career", "kind": "score", "request": make_request(
            "11:30", "12:30", declined_domains=["relationship", "career"],
        )},
        {"name": "refresh-probes-columns", "kind": "score", "request": make_request(
            "11:30", "12:30", refresh_probes=True, column_times=["11:40", "12:10", "12:20"],
        )},
        {"name": "minute-step-2", "kind": "score", "request": make_request("10:00", "13:00", minute_step=2)},
        {"name": "block-step-10", "kind": "block", "request": make_request("00:00", "23:59", minute_step=10)},
        {"name": "block-step-5", "kind": "block", "request": make_request("10:00", "14:00", minute_step=5)},
        {"name": "public-score-request", "kind": "score", "request": normalize_rectification_request(
            public_score_request(), today=FROZEN_TODAY,
        )},
        {"name": "diagnostics-11:00-13:00", "kind": "diagnostics", "request": make_request("11:00", "13:00")},
    ]
    return rows


def _strip_keys(kind: str) -> set[str]:
    # diagnostics no longer repeats the receipt. Drop it on both sides so the
    # remaining payload is what must stay byte-identical.
    if kind == "diagnostics":
        return {"column_compare_ms", "candidate_decision_receipt"}
    return {"column_compare_ms"}


class _SegmentTimer:
    def __init__(self) -> None:
        self.totals = {
            "guided_collect_windows": 0.0,
            "merge_transition_proximity": 0.0,
            "get_current_narayana_dasha": 0.0,
            "derive_functional_benefic_malefic": 0.0,
            "build_candidate_static_context": 0.0,
        }

    def wrap(self, owner: Any, name: str, key: str) -> Callable[[], None]:
        original = getattr(owner, name)

        def wrapped(*args: Any, **kwargs: Any) -> Any:
            started = time.perf_counter()
            try:
                return original(*args, **kwargs)
            finally:
                self.totals[key] += time.perf_counter() - started

        setattr(owner, name, wrapped)

        def restore() -> None:
            setattr(owner, name, original)

        return restore


def _median(samples: list[float]) -> float:
    return round(statistics.median(samples), 4)


def bench() -> dict[str, Any]:
    jobs = [
        ("12:00-12:20", "score", make_request("12:00", "12:20")),
        ("11:45-12:15", "score", make_request("11:45", "12:15")),
        ("11:00-13:00", "score", make_request("11:00", "13:00")),
        ("block-step-10", "block", make_request("00:00", "23:59", minute_step=10)),
    ]
    rows = []
    event_probes.date = _FrozenDate
    try:
        for name, kind, request in jobs:
            samples = []
            segments = []
            sizes = []
            counts = []
            for _ in range(5):
                timer = _SegmentTimer()
                restores = [
                    timer.wrap(event_probes, "guided_collect_windows", "guided_collect_windows"),
                    timer.wrap(scoring_service, "merge_transition_proximity", "merge_transition_proximity"),
                    timer.wrap(event_engine.narayana_dasha, "get_current_narayana_dasha", "get_current_narayana_dasha"),
                    timer.wrap(event_engine.functional_benefics, "derive_functional_benefic_malefic", "derive_functional_benefic_malefic"),
                    timer.wrap(event_engine, "build_candidate_static_context", "build_candidate_static_context"),
                ]
                started = time.perf_counter()
                try:
                    result = _call(kind, request)
                finally:
                    for restore in restores:
                        restore()
                elapsed = time.perf_counter() - started
                samples.append(elapsed)
                segments.append({key: round(value, 4) for key, value in timer.totals.items()})
                sizes.append(len(_canonical(_strip(result, _strip_keys(kind)))))
                counts.append(result.get("candidate_count") or len(result.get("candidate_scores") or []))
            median_segments: dict[str, float] = {}
            for key in segments[0]:
                median_segments[key] = _median([item[key] for item in segments])
            rows.append({
                "window": name,
                "candidate_count": counts[0],
                "median_seconds": _median(samples),
                "segment_median_seconds": median_segments,
                "response_bytes": int(statistics.median(sizes)),
            })
    finally:
        event_probes.date = date
    return {"runs": 5, "rows": rows}


def dump(out_dir: Path) -> list[dict[str, str]]:
    out_dir.mkdir(parents=True, exist_ok=True)
    written = []
    event_probes.date = _FrozenDate
    try:
        for row in scenarios():
            result = _call(row["kind"], row["request"], internal=True)
            payload = _strip(result, _strip_keys(row["kind"]))
            raw = _canonical(payload)
            digest = hashlib.sha256(raw).hexdigest()
            path = out_dir / f"{row['name'].replace(':', '')}.json"
            path.write_bytes(raw)
            written.append({"name": row["name"], "sha256": digest, "bytes": str(len(raw))})
            print(f"{row['name']}\t{digest}\t{len(raw)}")
    finally:
        event_probes.date = date
    return written


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("command", choices=("bench", "dump"))
    parser.add_argument("--out", type=Path, default=ROOT / "artifacts" / "rect-scoring-perf-20261013")
    args = parser.parse_args()
    if args.command == "bench":
        print(json.dumps(bench(), ensure_ascii=True, indent=2))
    else:
        dump(args.out)


if __name__ == "__main__":
    main()