#!/usr/bin/env python3
"""Summarize the non-performance CollectiveSpec P0 phase-trace artifacts."""

from __future__ import annotations

import argparse
import json
from collections import defaultdict
from pathlib import Path
from typing import Any


def load_json(path: Path) -> Any:
    return json.loads(path.read_text(encoding="utf-8"))


def integer(value: Any) -> int | None:
    return value if isinstance(value, int) and not isinstance(value, bool) else None


def merge_histograms(records: list[dict[str, Any]], field: str) -> dict[str, int]:
    result: dict[str, int] = {}
    for record in records:
        values = record.get(field)
        if not isinstance(values, dict):
            continue
        for key, value in values.items():
            count = integer(value)
            if count is not None and count >= 0:
                result[str(key)] = result.get(str(key), 0) + count
    return dict(sorted(result.items(), key=lambda item: int(item[0])))


def load_events(log_dir: Path) -> tuple[list[dict[str, Any]], list[str]]:
    events: list[dict[str, Any]] = []
    errors: list[str] = []
    for path in sorted(log_dir.glob("*.jsonl")):
        for line_number, raw in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1):
            if not raw:
                continue
            try:
                record = json.loads(raw)
            except json.JSONDecodeError as exc:
                errors.append(f"{path.name}:{line_number}: {exc.msg}")
                continue
            if not isinstance(record, dict):
                errors.append(f"{path.name}:{line_number}: event is not an object")
                continue
            events.append(record)
    return events, errors


def probe_integrity(cell: Path, expected_requests: int, expected_tokens: int) -> dict[str, Any]:
    details_paths = sorted(cell.glob("store/*/trials/trial-*/probe_details.jsonl"))
    result_paths = sorted(cell.glob("store/*/trials/trial-*/result.json"))
    if len(details_paths) != 1 or len(result_paths) != 1:
        return {
            "valid": False,
            "failures": [
                f"probe_details_count={len(details_paths)}",
                f"result_count={len(result_paths)}",
            ],
        }
    details_rows = [json.loads(line) for line in details_paths[0].read_text(encoding="utf-8").splitlines() if line]
    result = load_json(result_paths[0])
    failures: list[str] = []
    if result.get("status") != "completed":
        failures.append(f"result_status={result.get('status')}")
    if len(details_rows) != 1:
        failures.append(f"detail_rows={len(details_rows)}")
        return {"valid": False, "failures": failures}
    outcomes = details_rows[0].get("outcomes")
    if not isinstance(outcomes, list):
        failures.append("outcomes_not_list")
        outcomes = []
    if len(outcomes) != expected_requests:
        failures.append(f"outcome_count={len(outcomes)}_expected={expected_requests}")
    successful = sum(bool(item.get("success")) for item in outcomes if isinstance(item, dict))
    verified_tokens = sum(
        isinstance(item, dict)
        and item.get("completion_tokens_source") == "usage"
        and item.get("completion_tokens") == expected_tokens
        for item in outcomes
    )
    if successful != expected_requests:
        failures.append(f"success_count={successful}_expected={expected_requests}")
    if verified_tokens != expected_requests:
        failures.append(f"usage_token_count={verified_tokens}_expected={expected_requests}")
    return {
        "valid": not failures,
        "failures": failures,
        "outcome_count": len(outcomes),
        "success_count": successful,
        "usage_token_count": verified_tokens,
        "details_path": str(details_paths[0]),
        "result_path": str(result_paths[0]),
    }


def rank_key(event: dict[str, Any]) -> str | None:
    dp_rank = event.get("data_parallel_rank")
    local_rank = event.get("local_rank")
    if dp_rank is None or local_rank is None:
        return None
    return f"dp{dp_rank}/local{local_rank}"


def _event_order(event: dict[str, Any], epoch_field: str) -> tuple[int, float]:
    epoch = integer(event.get(epoch_field))
    monotonic_s = event.get("monotonic_s")
    return (
        epoch if epoch is not None else -1,
        float(monotonic_s) if isinstance(monotonic_s, (int, float)) else 0.0,
    )


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


def _target_event(event: dict[str, Any]) -> bool:
    """Exclude profile/dummy calls with an empty SchedulerOutput.

    The worker hook also wraps vLLM's profile and DP dummy runs. Those calls
    have a physical batch descriptor but no target request, so their ranks and
    counts are not a valid collective-plan invariant.
    """

    return (
        (integer(event.get("request_count")) or 0) > 0
        and (integer(event.get("total_scheduled_rows")) or 0) > 0
    )


def _target_signature(event: dict[str, Any]) -> str:
    return _canonical(
        {
            "cudagraph_mode": event.get("cudagraph_mode"),
            "ordered_plan_digest": event.get("ordered_plan_digest"),
            "physical_batch_rows": event.get("physical_batch_rows"),
            "rows_across_dp": event.get("rows_across_dp"),
            "scheduled_dp_metadata": event.get("scheduled_dp_metadata"),
            "semantic_plan_digest": event.get("semantic_plan_digest"),
            "should_ubatch": event.get("should_ubatch"),
            "total_scheduled_rows": event.get("total_scheduled_rows"),
        }
    )


def _target_metadata_check(event: dict[str, Any]) -> tuple[bool, int | None, bool | None]:
    """Check the local, post-DP-coordination shape invariants."""

    metadata = event.get("scheduled_dp_metadata")
    dp_rank = integer(event.get("data_parallel_rank"))
    dp_size = integer(event.get("data_parallel_size"))
    total_rows = integer(event.get("total_scheduled_rows"))
    physical_rows = integer(event.get("physical_batch_rows"))
    rows_across_dp = event.get("rows_across_dp")
    if (
        not isinstance(metadata, dict)
        or dp_rank is None
        or dp_size is None
        or dp_rank < 0
        or dp_rank >= dp_size
        or total_rows is None
        or physical_rows is None
        or not isinstance(rows_across_dp, list)
    ):
        return False, None, None
    raw_rows = metadata.get("num_tokens_per_rank")
    padded_rows = metadata.get("num_tokens_across_dp")
    if (
        not isinstance(raw_rows, list)
        or not isinstance(padded_rows, list)
        or len(raw_rows) != dp_size
        or len(padded_rows) != dp_size
        or len(rows_across_dp) != dp_size
        or not all(integer(value) is not None and integer(value) >= 0 for value in raw_rows)
        or not all(integer(value) is not None and integer(value) >= 0 for value in padded_rows)
        or not all(integer(value) is not None and integer(value) >= 0 for value in rows_across_dp)
    ):
        return False, None, None
    raw = [int(value) for value in raw_rows]
    padded = [int(value) for value in padded_rows]
    coordinated = [int(value) for value in rows_across_dp]
    valid = (
        raw[dp_rank] == total_rows
        and physical_rows == coordinated[dp_rank]
        and padded == coordinated
        and all(padded_value >= raw_value for padded_value, raw_value in zip(padded, raw))
    )
    return valid, padded[dp_rank] - raw[dp_rank], len(set(raw)) > 1


def _dp_pair_key(event: dict[str, Any]) -> tuple[str, str] | None:
    """Return the TP-shard key shared by the two actual DP peers.

    New P0 logs carry ``tensor_parallel_rank``. The v2 artifact predates that
    field, where ``Worker.rank`` was logged as ``global_rank`` and is the same
    TP-shard identifier within each DP replica. The fallback is explicit in
    the output so it cannot be confused with a true global rank.
    """

    tp_rank = integer(event.get("tensor_parallel_rank"))
    if tp_rank is not None:
        return f"tp{tp_rank}", "tensor_parallel_rank"
    legacy_worker_rank = integer(event.get("global_rank"))
    if legacy_worker_rank is not None:
        return f"legacy-worker-rank{legacy_worker_rank}", "legacy_worker_rank"
    local_rank = integer(event.get("local_rank"))
    tp_size = integer(event.get("tensor_parallel_size"))
    if local_rank is not None and tp_size is not None and tp_size > 0:
        return f"local-rank-mod-tp{local_rank % tp_size}", "local_rank_mod_tp_size"
    return None


def _dp_coordination_signature(event: dict[str, Any]) -> str:
    metadata = event.get("scheduled_dp_metadata")
    if not isinstance(metadata, dict):
        metadata = {}
    return _canonical(
        {
            "all_decode": metadata.get("all_decode"),
            "all_prefill": metadata.get("all_prefill"),
            "cudagraph_mode": metadata.get("cudagraph_mode"),
            "input_fits_in_drafter": metadata.get("input_fits_in_drafter"),
            "is_prompt_batch": metadata.get("is_prompt_batch"),
            "num_reqs_per_rank": metadata.get("num_reqs_per_rank"),
            "num_tokens_across_dp": metadata.get("num_tokens_across_dp"),
            "num_tokens_per_rank": metadata.get("num_tokens_per_rank"),
            "rows_across_dp": event.get("rows_across_dp"),
            "should_ubatch": metadata.get("should_ubatch"),
        }
    )


def dp_pair_coordination(target_events: list[dict[str, Any]]) -> dict[str, Any]:
    """Compare scalar DP coordination seen by the two peers of each TP shard."""

    grouped: dict[str, dict[int, dict[int, dict[str, Any]]]] = defaultdict(
        lambda: defaultdict(dict)
    )
    sources: dict[str, set[str]] = defaultdict(set)
    expected_sizes: dict[str, set[int]] = defaultdict(set)
    for event in target_events:
        key = _dp_pair_key(event)
        dp_rank = integer(event.get("data_parallel_rank"))
        epoch = integer(event.get("batch_phase_epoch"))
        dp_size = integer(event.get("data_parallel_size"))
        if key is None or dp_rank is None or epoch is None:
            continue
        group, source = key
        grouped[group][dp_rank][epoch] = event
        sources[group].add(source)
        if dp_size is not None:
            expected_sizes[group].add(dp_size)

    groups: dict[str, Any] = {}
    for group, by_dp_rank in grouped.items():
        dp_ranks = sorted(by_dp_rank)
        epoch_sets = [set(by_epoch) for by_epoch in by_dp_rank.values()]
        shared_epochs = sorted(set.intersection(*epoch_sets)) if epoch_sets else []
        signatures = {
            dp_rank: {
                epoch: _dp_coordination_signature(event)
                for epoch, event in by_epoch.items()
            }
            for dp_rank, by_epoch in by_dp_rank.items()
        }
        shared_equal = all(
            len({signatures[dp_rank][epoch] for dp_rank in dp_ranks}) == 1
            for epoch in shared_epochs
        )
        expected_values = expected_sizes.get(group, set())
        expected_dp_size = next(iter(expected_values)) if len(expected_values) == 1 else None
        groups[group] = {
            "grouping_source": sorted(sources[group]),
            "data_parallel_ranks": dp_ranks,
            "expected_data_parallel_size": expected_dp_size,
            "target_epoch_count_by_dp_rank": {
                str(dp_rank): len(by_dp_rank[dp_rank]) for dp_rank in dp_ranks
            },
            "shared_target_epoch_count": len(shared_epochs),
            "all_shared_epoch_signatures_identical": bool(shared_epochs) and shared_equal,
            "all_expected_dp_peers_observed": expected_dp_size is not None
            and len(dp_ranks) == expected_dp_size,
        }
    return {
        "group_count": len(groups),
        "groups": groups,
        "all_expected_dp_peers_observed": bool(groups)
        and all(value["all_expected_dp_peers_observed"] for value in groups.values()),
        "all_shared_epoch_signatures_identical": bool(groups)
        and all(value["all_shared_epoch_signatures_identical"] for value in groups.values()),
    }


def target_execution_agreement(worker_events: list[dict[str, Any]]) -> dict[str, Any]:
    batch_events = [event for event in worker_events if event.get("event") == "batch_execution_plan"]
    target_events = [event for event in batch_events if _target_event(event)]
    per_rank: dict[str, list[dict[str, Any]]] = defaultdict(list)
    per_replica: dict[str, set[str]] = defaultdict(set)
    for event in target_events:
        rank = rank_key(event)
        dp_rank = integer(event.get("data_parallel_rank"))
        if rank is None or dp_rank is None:
            continue
        per_rank[rank].append(event)
        per_replica[str(dp_rank)].add(rank)

    sequences: dict[str, list[str]] = {}
    for rank, values in per_rank.items():
        values.sort(key=lambda event: _event_order(event, "batch_phase_epoch"))
        sequences[rank] = [_target_signature(event) for event in values]
    counts = {rank: len(values) for rank, values in sequences.items()}

    replicas: dict[str, Any] = {}
    for dp_rank, ranks in per_replica.items():
        rank_list = sorted(ranks)
        distinct = {tuple(sequences[rank]) for rank in rank_list}
        phase_counts = [counts[rank] for rank in rank_list]
        replicas[dp_rank] = {
            "ranks": rank_list,
            "phase_count_by_rank": {rank: counts[rank] for rank in rank_list},
            "phase_count_equal": bool(phase_counts) and len(set(phase_counts)) == 1,
            "sequence_distinct_count": len(distinct),
            "identical_target_execution_sequence": len(distinct) == 1,
        }

    metadata_valid_count = 0
    anchor_local_ranks: dict[int, int] = {}
    for event in target_events:
        dp_rank = integer(event.get("data_parallel_rank"))
        local_rank = integer(event.get("local_rank"))
        if dp_rank is None or local_rank is None:
            continue
        previous = anchor_local_ranks.get(dp_rank)
        if previous is None or local_rank < previous:
            anchor_local_ranks[dp_rank] = local_rank
    anchor_events = [
        event
        for event in target_events
        if integer(event.get("data_parallel_rank")) in anchor_local_ranks
        and integer(event.get("local_rank"))
        == anchor_local_ranks[integer(event.get("data_parallel_rank"))]
    ]
    metadata_padding: list[int] = []
    raw_unequal_count = 0
    for event in target_events:
        valid, _, _ = _target_metadata_check(event)
        metadata_valid_count += valid
    for event in anchor_events:
        _, padding, raw_unequal = _target_metadata_check(event)
        if padding is not None:
            metadata_padding.append(padding)
        raw_unequal_count += bool(raw_unequal)
    return {
        "batch_execution_plan_count": len(batch_events),
        "dummy_or_profile_event_count": len(batch_events) - len(target_events),
        "target_event_count": len(target_events),
        "worker_rank_count": len(sequences),
        "target_phase_count_by_rank": counts,
        "within_dp_replica_tp_group": replicas,
        "all_target_phase_counts_equal_within_dp_replica": bool(replicas)
        and all(value["phase_count_equal"] for value in replicas.values()),
        "all_target_sequences_identical_within_dp_replica": bool(replicas)
        and all(value["identical_target_execution_sequence"] for value in replicas.values()),
        "dp_metadata": {
            "target_record_count": len(target_events),
            "valid_record_count": metadata_valid_count,
            "all_records_valid": len(target_events) > 0
            and metadata_valid_count == len(target_events),
            "anchor_record_count": len(anchor_events),
            "raw_dp_counts_unequal_anchor_record_count": raw_unequal_count,
            "per_dp_replica_anchor_post_coordinate_minus_local_raw_rows": {
                "n": len(metadata_padding),
                "min": min(metadata_padding) if metadata_padding else None,
                "max": max(metadata_padding) if metadata_padding else None,
                "sum": sum(metadata_padding),
                "distinct_value_count": len(set(metadata_padding)),
                "non_negative": all(value >= 0 for value in metadata_padding),
            },
        },
        "dp_pair_coordination": dp_pair_coordination(target_events),
    }


def summarize_cell(cell: Path, expected_requests: int, expected_tokens: int) -> dict[str, Any]:
    events, parse_errors = load_events(cell / "p0_logs")
    scheduler_events = [event for event in events if event.get("role") == "scheduler"]
    worker_events = [event for event in events if event.get("role") == "worker"]
    candidate_events = [event for event in scheduler_events if event.get("event") == "candidate_truncate"]
    schedule_events = [event for event in scheduler_events if event.get("event") == "schedule"]
    verify_schedule_events = [
        event for event in schedule_events if integer(event.get("spec_request_count")) not in (None, 0)
    ]
    batch_events = [event for event in worker_events if event.get("event") == "batch_execution_plan"]
    target_events = [event for event in batch_events if _target_event(event)]
    rows_across_dp = [event.get("rows_across_dp") for event in target_events]
    physical_rows = [
        integer(event.get("physical_batch_rows"))
        for event in target_events
        if integer(event.get("physical_batch_rows")) is not None
    ]
    return {
        "probe_integrity": probe_integrity(cell, expected_requests, expected_tokens),
        "event_parse_errors": parse_errors,
        "event_count": len(events),
        "scheduler": {
            "schedule_event_count": len(schedule_events),
            "verify_schedule_event_count": len(verify_schedule_events),
            "candidate_event_count": len(candidate_events),
            "before_k_histogram": merge_histograms(candidate_events, "before_k_histogram"),
            "after_k_histogram": merge_histograms(candidate_events, "after_k_histogram"),
            "after_k_distinct_value_count": len(merge_histograms(candidate_events, "after_k_histogram")),
            "dp_rank_count": len(
                {
                    event.get("data_parallel_rank")
                    for event in scheduler_events
                    if event.get("data_parallel_rank") is not None
                }
            ),
        },
        "worker": {
            **target_execution_agreement(worker_events),
            "target_rows_across_dp_present_count": sum(value is not None for value in rows_across_dp),
            "target_physical_batch_rows": {
                "n": len(physical_rows),
                "min": min(physical_rows) if physical_rows else None,
                "max": max(physical_rows) if physical_rows else None,
                "distinct_value_count": len(set(physical_rows)),
                "non_negative": all(value >= 0 for value in physical_rows),
            },
        },
    }


def data_sanity(cells: dict[str, dict[str, Any]], expected_gpu_count: int | None) -> dict[str, Any]:
    event_counts = [integer(cell.get("event_count")) or 0 for cell in cells.values()]
    candidate_distinct = [
        integer(cell.get("scheduler", {}).get("after_k_distinct_value_count")) or 0
        for cell in cells.values()
    ]
    return {
        "n_cells": len(cells),
        "event_count": {
            "n": len(event_counts),
            "min": min(event_counts) if event_counts else None,
            "max": max(event_counts) if event_counts else None,
            "distinct_value_count": len(set(event_counts)),
            "non_negative": all(value >= 0 for value in event_counts),
        },
        "candidate_k_distinct": {
            "n": len(candidate_distinct),
            "min": min(candidate_distinct) if candidate_distinct else None,
            "max": max(candidate_distinct) if candidate_distinct else None,
            "distinct_value_count": len(set(candidate_distinct)),
            "non_negative": all(value >= 0 for value in candidate_distinct),
        },
        "invariants": {
            "all_probe_integrity_valid": all(
                bool(cell.get("probe_integrity", {}).get("valid")) for cell in cells.values()
            ),
            "no_event_parse_errors": all(not cell.get("event_parse_errors") for cell in cells.values()),
            "all_target_phase_counts_equal_within_dp_replica": all(
                bool(
                    cell.get("worker", {}).get(
                        "all_target_phase_counts_equal_within_dp_replica"
                    )
                )
                for cell in cells.values()
            ),
            "all_target_sequences_identical_within_dp_replica": all(
                bool(
                    cell.get("worker", {}).get(
                        "all_target_sequences_identical_within_dp_replica"
                    )
                )
                for cell in cells.values()
            ),
            "all_target_dp_metadata_valid": all(
                bool(cell.get("worker", {}).get("dp_metadata", {}).get("all_records_valid"))
                for cell in cells.values()
            ),
            "all_target_dp_coordination_records_observed": all(
                integer(cell.get("worker", {}).get("target_rows_across_dp_present_count"))
                == integer(cell.get("worker", {}).get("target_event_count"))
                and integer(cell.get("worker", {}).get("target_event_count", 0)) > 0
                for cell in cells.values()
            ),
            "all_expected_dp_peers_observed": all(
                bool(
                    cell.get("worker", {}).get("dp_pair_coordination", {}).get(
                        "all_expected_dp_peers_observed"
                    )
                )
                for cell in cells.values()
            ),
            "all_shared_dp_pair_coordination_signatures_identical": all(
                bool(
                    cell.get("worker", {}).get("dp_pair_coordination", {}).get(
                        "all_shared_epoch_signatures_identical"
                    )
                )
                for cell in cells.values()
            ),
            "all_expected_workers_observed": (
                all(
                    integer(cell.get("worker", {}).get("worker_rank_count")) == expected_gpu_count
                    for cell in cells.values()
                )
                if expected_gpu_count is not None
                else None
            ),
        },
    }


def markdown(payload: dict[str, Any]) -> str:
    lines = [
        "# CollectiveSpec P0 phase-trace summary",
        "",
        "P0 verifies the collective/liveness premise only; it is not a performance result.",
        "",
        "| policy | complete usage-verified requests | candidate K values observed | worker ranks | target phase counts equal within each DP replica | target sequences identical within each DP replica |",
        "|---|---:|---:|---:|---:|---:|",
    ]
    for name, cell in payload["cells"].items():
        probe = cell["probe_integrity"]
        scheduler = cell["scheduler"]
        worker = cell["worker"]
        values = ",".join(sorted(scheduler["after_k_histogram"])) or "—"
        lines.append(
            "| {name} | {success}/{count} ({valid}) | {values} | {ranks} | {equal} | {within} |".format(
                name=name,
                success=probe.get("success_count", "—"),
                count=probe.get("outcome_count", "—"),
                valid=probe.get("valid", False),
                values=values,
                ranks=worker.get("worker_rank_count", 0),
                equal=worker.get("all_target_phase_counts_equal_within_dp_replica", False),
                within=worker.get("all_target_sequences_identical_within_dp_replica", False),
            )
        )
    lines.extend(
        [
            "",
            "## Data sanity",
            "",
            "```json",
            json.dumps(payload["data_sanity"], indent=2, sort_keys=True),
            "```",
            "",
        ]
    )
    return "\n".join(lines)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--root", type=Path, required=True)
    parser.add_argument("--output-json", type=Path, required=True)
    parser.add_argument("--output-md", type=Path, required=True)
    args = parser.parse_args()
    manifest = load_json(args.root / "manifest.json")
    expected_requests = integer(manifest.get("request_count"))
    expected_tokens = integer(manifest.get("completion_tokens"))
    if expected_requests is None or expected_tokens is None:
        raise SystemExit("manifest must contain integer request_count and completion_tokens")
    policies = manifest.get("policies")
    if not isinstance(policies, list) or not all(isinstance(value, str) for value in policies):
        raise SystemExit("manifest policies must be a list of strings")
    cells = {
        policy: summarize_cell(args.root / policy, expected_requests, expected_tokens)
        for policy in policies
    }
    payload = {
        "manifest": manifest,
        "cells": cells,
        "data_sanity": data_sanity(cells, integer(manifest.get("gpu_count"))),
    }
    args.output_json.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    args.output_md.write_text(markdown(payload), encoding="utf-8")
    print(json.dumps(payload, indent=2, sort_keys=True))
    return 0 if all(payload["data_sanity"]["invariants"].values()) else 2


if __name__ == "__main__":
    raise SystemExit(main())
