"""Summarize exported xctrace intervals for the attached target only.

Reads XML incrementally; never prints the TOC environment or other processes.
GPU intervals are not shader timings. Nested intervals are not added to level 0.
Per-buffer covered time merges overlaps; group sums are not GPU utilization.
"""

import argparse
import io
import json
import math
import re
import statistics
import sys
import xml.etree.ElementTree as ET


def summarize(source, schemas, pid):
    ids, totals, buffers = {}, {}, {}
    scalar_tags = {"start-time", "duration", "process", "gpu-channel-name",
                   "gpu-state", "metal-nesting-level", "metal-command-buffer-id"}
    schema = None

    def value(element):
        if element is None:
            return None
        if "ref" in element.attrib:
            return ids[element.attrib["ref"]]
        return element.get("fmt") if element.tag == "process" else element.text

    for event, element in ET.iterparse(source, events=("start", "end")):
        if event == "start":
            if element.tag == "node":
                index = int(re.search(r"table\[(\d+)\]$", element.attrib["xpath"])[1])
                schema = schemas[index - 1]
            continue
        if element.tag in scalar_tags and "id" in element.attrib:
            ids[element.attrib["id"]] = value(element)
        if element.tag == "row":
            process = value(element.find("process"))
            if process and process.endswith(f"({pid})"):
                duration = value(element.find("duration"))
                if duration is not None:
                    ms = int(duration) / 1e6
                    channel = value(element.find("gpu-channel-name"))
                    state = value(element.find("gpu-state"))
                    depth = value(element.find("metal-nesting-level"))
                    key = "/".join(str(x) for x in (schema, channel, state, depth))
                    total = totals.setdefault(key, {"rows": 0, "sum_ms": 0, "max_ms": 0})
                    total["rows"] += 1
                    total["sum_ms"] += ms
                    total["max_ms"] = max(total["max_ms"], ms)
                    if (schema == "metal-gpu-intervals" and channel == "Compute"
                            and depth == "0" and state == "Active"):
                        command = value(element.find("metal-command-buffer-id"))
                        start = int(value(element.find("start-time"))) / 1e6
                        record = buffers.setdefault(command, {"command_buffer": command,
                            "start_ms": start, "end_ms": start, "gpu_interval_sum_ms": 0,
                            "max_interval_ms": 0, "intervals": 0, "_spans": []})
                        record["start_ms"] = min(record["start_ms"], start)
                        record["end_ms"] = max(record["end_ms"], start + ms)
                        record["gpu_interval_sum_ms"] += ms
                        record["max_interval_ms"] = max(record["max_interval_ms"], ms)
                        record["intervals"] += 1
                        record["_spans"].append((start, start + ms))
            element.clear()
        elif element.tag == "node":
            element.clear()
    for record in buffers.values():
        covered, end = 0, float("-inf")
        for start, stop in sorted(record.pop("_spans")):
            covered += max(0, stop - max(start, end))
            end = max(end, stop)
        record["gpu_covered_ms"] = covered
    return {"pid": pid, "interval_groups": totals,
            "largest_compute_buffers": sorted(buffers.values(),
                key=lambda x: x["gpu_interval_sum_ms"], reverse=True)[:12]}


def self_test():
    xml = """<trace-query-result><node xpath="//table[1]">
    <row><start-time id="1">1000000</start-time><duration id="2">2000000</duration>
    <process id="3" fmt="ds4-server (42)"/><gpu-channel-name id="4">Compute</gpu-channel-name>
    <gpu-state id="5">Active</gpu-state><metal-nesting-level id="6">0</metal-nesting-level>
    <metal-command-buffer-id id="7">100</metal-command-buffer-id></row>
    <row><start-time>4000000</start-time><duration ref="2"/><process ref="3"/>
    <gpu-channel-name ref="4"/><gpu-state ref="5"/><metal-nesting-level ref="6"/>
    <metal-command-buffer-id ref="7"/></row>
    <row><start-time>2000000</start-time><duration ref="2"/><process ref="3"/>
    <gpu-channel-name ref="4"/><gpu-state ref="5"/><metal-nesting-level ref="6"/>
    <metal-command-buffer-id ref="7"/></row>
    <row><start-time>1000000</start-time><duration ref="2"/><process ref="3"/>
    <gpu-channel-name ref="4"/><gpu-state ref="5"/><metal-nesting-level>1</metal-nesting-level>
    <metal-command-buffer-id ref="7"/></row>
    <row><duration ref="2"/><process fmt="other (43)"/></row>
    </node></trace-query-result>"""
    result = summarize(io.StringIO(xml), ["metal-gpu-intervals"], 42)
    assert list(result["interval_groups"].values()) == [
        {"rows": 3, "sum_ms": 6, "max_ms": 2}, {"rows": 1, "sum_ms": 2, "max_ms": 2}]
    assert result["largest_compute_buffers"] == [{"command_buffer": "100", "start_ms": 1,
        "end_ms": 6, "gpu_interval_sum_ms": 6, "gpu_covered_ms": 5,
        "max_interval_ms": 2, "intervals": 3}]
    print("PASS: XML references, PID filtering, depth filtering and overlap union")


def summarize_stages(source):
    """Validate native counter receipts; group encoder spans, never shaders."""
    headers, stages, groups = {}, {}, {}
    for line in source:
        # Supervisor JSON can precede a worker's complete native stderr record.
        # Keep strict field/completeness checks; never discard a damaged record.
        _, marker, record = line.partition("ds4: stage-counter ")
        if not marker:
            continue
        values = dict(word.split("=", 1) for word in record.split())
        cb = int(values["cb"])
        if "stages" in values:
            assert cb not in headers, f"duplicate CB {cb}"
            assert values["resolved"] == "1" and values["dropped"] == "0", values
            headers[cb] = values
            stages[cb] = {}
            continue
        assert cb in headers and values["valid"] == "1", values
        index = int(values["stage"])
        assert index not in stages[cb], (cb, index)
        start, end = int(values["start"]), int(values["end"])
        duration = float(values["ms"])
        scale = float(headers[cb]["ns_per_tick"])
        assert end >= start and math.isfinite(duration) and duration >= 0
        assert abs(duration - (end - start) * scale / 1e6) <= 0.000002
        stages[cb][index] = (start, end)
        key = (values["first"], values["last"], int(values["qwen_dispatches"]))
        groups.setdefault(key, []).append(duration)
    assert headers, "no counter receipts"
    covered_total, cb_total = 0.0, 0.0
    for cb, header in headers.items():
        assert set(stages[cb]) == set(range(int(header["stages"]))), cb
        covered, prior_end = 0, 0
        for start, end in sorted(stages[cb].values()):
            covered += max(0, end - max(start, prior_end))
            prior_end = max(prior_end, end)
        covered_ms = covered * float(header["ns_per_tick"]) / 1e6
        cb_ms = float(header["cb_ms"])
        assert math.isfinite(cb_ms) and cb_ms > 0
        assert covered_ms <= cb_ms * 1.05 + 0.002, (cb, covered_ms, cb_ms)
        covered_total += covered_ms
        cb_total += cb_ms
    return {"scope": "Qwen four-row command buffers, existing compute encoder boundaries",
            "command_buffers": len(headers), "encoder_spans": sum(map(len, stages.values())),
            "counter_covered_ms": covered_total, "command_buffer_gpu_ms": cb_total,
            "groups_are_encoder_spans_not_individual_kernels": True,
            "encoder_groups": sorted([
                {"first_qwen_kernel": first, "last_qwen_kernel": last,
                 "qwen_dispatches": dispatches, "samples": len(values),
                 "sum_ms": sum(values), "median_ms": statistics.median(values), "max_ms": max(values)}
                for (first, last, dispatches), values in groups.items()
            ], key=lambda row: row["sum_ms"], reverse=True)}


def stage_self_test():
    log = ("ds4: stage-counter cb=1 rows=4 stages=2 dropped=0 resolved=1 ns_per_tick=1 cb_ms=0.05\n"
           "ds4: stage-counter cb=1 stage=0 valid=1 start=100000 end=130000 ms=0.03 qwen_dispatches=2 first=a last=b\n"
           "ds4: stage-counter cb=1 stage=1 valid=1 start=120000 end=140000 ms=0.02 qwen_dispatches=1 first=c last=c\n")
    result = summarize_stages(io.StringIO(log))
    assert result["encoder_spans"] == 2 and result["counter_covered_ms"] == 0.04
    assert summarize_stages(io.StringIO(log.replace("ds4:", '{"event":ds4:'))) == result
    for bad in (log.replace("resolved=1", "resolved=0"), log.replace("dropped=0", "dropped=1"),
                log.replace("ms=0.03", "ms=1.25"), "\n".join(log.splitlines()[:-1])):
        try:
            summarize_stages(io.StringIO(bad))
        except AssertionError:
            continue
        raise AssertionError("invalid stage trace accepted")
    print("PASS: stage completeness, failure reporting, clock scale and overlap coverage")


if __name__ == "__main__":
    if sys.argv[1:] == ["--self-test"]:
        self_test()
        stage_self_test()
    else:
        parser = argparse.ArgumentParser(description=__doc__)
        parser.add_argument("toc", nargs="?")
        parser.add_argument("intervals", nargs="?")
        parser.add_argument("--stage-log", help="Native DS4_METAL_STAGE_PROFILE stderr receipt")
        args = parser.parse_args()
        if args.stage_log:
            if args.toc or args.intervals:
                parser.error("do not combine --stage-log with XML inputs")
            with open(args.stage_log) as source:
                print(json.dumps(summarize_stages(source), indent=2))
            sys.exit(0)
        if not args.toc or not args.intervals:
            parser.error("provide both TOC and intervals, or --stage-log")
        toc = ET.parse(args.toc)
        if len(toc.findall("./run")) != 1:
            parser.error("export a trace containing exactly one run")
        pid = int(toc.find("./run/info/target/process").attrib["pid"])
        schemas = [table.attrib["schema"] for table in toc.findall("./run/data/table")]
        result = summarize(args.intervals, schemas, pid)
        result["trace_duration_seconds"] = float(toc.findtext("./run/info/summary/duration"))
        print(json.dumps(result, indent=2))