#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import re
import struct
import sys
from collections import Counter, namedtuple
from pathlib import Path


Event = namedtuple("Event", "op peer size ordinal instruction")
Program = namedtuple("Program", "events operations starts_inactive ends_with_send")
Transfer = namedtuple(
    "Transfer",
    "sender receiver size sender_ordinal receiver_ordinal sender_instruction receiver_instruction "
    "sender_last_instruction receiver_last_instruction count",
)
CORE_FILE = re.compile(r"core_(\d+)\.(json|pim)$")
HEADER = struct.Struct("<4sII")
RECORD = struct.Struct("<BBBBiiii")
OPCODE_NAMES = (
    "nop", "sldi", "sld", "sadd", "ssub", "smul", "saddi", "smuli", "setbw",
    "mvmul", "vvadd", "vvsub", "vvmul", "vvdmul", "vvmax", "vvsll", "vvsra",
    "vavg", "vrelu", "vtanh", "vsigm", "vsoftmax", "vmv", "vrsu", "vrsl",
    "ld", "st", "lldi", "lmv", "send", "recv", "wait", "sync",
)
IGNORED_OPS = {"nop", "sldi", "lldi", "setbw"}
MEMORY_OPS = {"sld", "ld", "st", "vmv", "vrsu", "vrsl", "lmv"}
DISPLAY_ORDER = (
    "ld", "lmv", "st", "sld", "vmv", "vrsu", "vrsl",
    "vvmul", "mvmul", "vvdmul", "smul", "smuli",
    "vvadd", "vvsub", "vvmax", "vavg", "sadd", "saddi", "ssub",
    "vvsll", "vvsra", "vrelu", "vtanh", "vsigm", "vsoftmax", "wait", "sync",
)
DISPLAY_RANK = {op: rank for rank, op in enumerate(DISPLAY_ORDER)}
ARROW_HEIGHT = 2.0


def summarize_operations(operations: tuple[str, ...] | list[str]) -> str:
    counts = Counter(op for op in operations if op not in IGNORED_OPS and op not in ("send", "recv"))
    ordered = sorted(counts, key=lambda op: (DISPLAY_RANK.get(op, len(DISPLAY_RANK)), op))
    groups = ([op for op in ordered if op in MEMORY_OPS], [op for op in ordered if op not in MEMORY_OPS])
    rows = [
        [(counts[op], op) for op in group[index : index + 2]]
        for group in groups
        for index in range(0, len(group), 2)
    ]
    if not rows:
        return ""
    count_width = max(len(str(count)) for row in rows for count, _ in row)
    op_width = max(len(op) for row in rows for _, op in row)
    return '""' + "\\n".join(
        "  ".join(f"{count:>{count_width}}x {op:<{op_width}}" for count, op in row).rstrip()
        for row in rows
    ) + '""'


def read_json(path: Path) -> Program:
    instructions = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(instructions, list):
        raise ValueError(f"{path}: expected a JSON instruction array")

    operations = tuple(
        str(instruction.get("op", "unknown")) if isinstance(instruction, dict) else "unknown"
        for instruction in instructions
    )
    events = []
    for index, instruction in enumerate(instructions):
        if not isinstance(instruction, dict):
            continue
        op = instruction.get("op")
        if op not in ("send", "recv"):
            continue
        try:
            peer = int(instruction["core"])
            size = int(instruction["size"])
        except (KeyError, TypeError, ValueError) as error:
            raise ValueError(f"{path}: invalid communication instruction {instruction!r}") from error
        events.append(Event(op, peer, size, len(events), index))
    starts_inactive = bool(
        events
        and events[0].op == "recv"
        and all(
            isinstance(instruction, dict) and instruction.get("op") == "sldi"
            for instruction in instructions[:events[0].instruction]
        )
    )
    last_op = instructions[-1].get("op") if instructions and isinstance(instructions[-1], dict) else None
    return Program(events, operations, starts_inactive, last_op == "send")


def read_binary(path: Path) -> Program:
    data = path.read_bytes()
    if len(data) < HEADER.size:
        raise ValueError(f"{path}: binary core file is too small")
    magic, version, count = HEADER.unpack_from(data)
    expected_size = HEADER.size + count * RECORD.size
    if magic != b"PIMB":
        raise ValueError(f"{path}: invalid PIM binary magic")
    if version != 1:
        raise ValueError(f"{path}: unsupported PIM binary version {version}")
    if len(data) != expected_size:
        raise ValueError(f"{path}: expected {expected_size} bytes, found {len(data)}")

    events = []
    operations = []
    last_opcode = None
    for index in range(count):
        opcode, _, _, _, peer, _, _, size = RECORD.unpack_from(data, HEADER.size + index * RECORD.size)
        last_opcode = opcode
        operations.append(OPCODE_NAMES[opcode] if opcode < len(OPCODE_NAMES) else f"opcode_{opcode}")
        if opcode in (29, 30):
            events.append(Event("send" if opcode == 29 else "recv", peer, size, len(events), index))
    starts_inactive = bool(
        events
        and events[0].op == "recv"
        and all(data[HEADER.size + index * RECORD.size] == 1 for index in range(events[0].instruction))
    )
    return Program(events, tuple(operations), starts_inactive, last_opcode == 29)


def read_programs(directory: Path, artifact_format: str) -> dict[int, Program]:
    candidates: dict[str, dict[int, Path]] = {"json": {}, "pim": {}}
    for path in directory.iterdir():
        match = CORE_FILE.fullmatch(path.name)
        if match:
            candidates[match.group(2)][int(match.group(1))] = path

    if artifact_format == "auto":
        artifact_format = "json" if candidates["json"] else "pim"
    files = candidates[artifact_format]
    if not files:
        raise ValueError(f"{directory}: no core_*.{artifact_format} files found")

    reader = read_json if artifact_format == "json" else read_binary
    return {core: reader(path) for core, path in files.items()}


def match_transfers(programs: dict[int, Program]) -> list[Transfer]:
    positions = {core: 0 for core in programs}
    transfers = []

    while True:
        made_progress = False
        for core in sorted(programs):
            events = programs[core].events
            position = positions[core]
            if position == len(events):
                continue
            event = events[position]

            if event.peer not in programs:
                sender, receiver = (core, event.peer) if event.op == "send" else (event.peer, core)
                transfers.append(
                    Transfer(
                        sender,
                        receiver,
                        event.size,
                        event.ordinal if event.op == "send" else None,
                        event.ordinal if event.op == "recv" else None,
                        event.instruction if event.op == "send" else None,
                        event.instruction if event.op == "recv" else None,
                        event.instruction if event.op == "send" else None,
                        event.instruction if event.op == "recv" else None,
                        1,
                    )
                )
                positions[core] += 1
                made_progress = True
                continue

            peer_events = programs[event.peer].events
            peer_position = positions[event.peer]
            if peer_position == len(peer_events):
                continue
            peer_event = peer_events[peer_position]
            if (
                event.peer == core
                or peer_event.peer != core
                or peer_event.op == event.op
                or peer_event.size != event.size
            ):
                continue

            send = event if event.op == "send" else peer_event
            receive = peer_event if event.op == "send" else event
            sender, receiver = (core, event.peer) if event.op == "send" else (event.peer, core)
            transfers.append(
                Transfer(
                    sender,
                    receiver,
                    event.size,
                    send.ordinal,
                    receive.ordinal,
                    send.instruction,
                    receive.instruction,
                    send.instruction,
                    receive.instruction,
                    1,
                )
            )
            positions[core] += 1
            positions[event.peer] += 1
            made_progress = True

        if made_progress:
            continue
        remaining = {
            core: programs[core].events[position]
            for core, position in positions.items()
            if position < len(programs[core].events)
        }
        if not remaining:
            return transfers
        details = ", ".join(
            f"core {core}: {event.op} {event.peer} ({event.size} B)"
            for core, event in sorted(remaining.items())
        )
        raise ValueError(f"communication streams cannot be matched at {details}")


def visible_transfers(
    transfers: list[Transfer],
    selected: set[int],
    programs: dict[int, Program],
) -> list[Transfer]:
    visible = [transfer for transfer in transfers if transfer.sender in selected or transfer.receiver in selected]
    grouped = []
    for transfer in visible:
        previous = grouped[-1] if grouped else None
        can_group = (
            previous is not None
            and transfer.sender_ordinal is not None
            and previous.sender == transfer.sender
            and previous.receiver == transfer.receiver
            and previous.sender_ordinal + previous.count == transfer.sender_ordinal
            and (
                transfer.receiver_ordinal is None
                or previous.receiver_ordinal + previous.count == transfer.receiver_ordinal
            )
        )
        if can_group:
            sender_gap = programs[transfer.sender].operations[
                previous.sender_last_instruction + 1 : transfer.sender_instruction
            ]
            receiver_gap = (
                programs[transfer.receiver].operations[
                    previous.receiver_last_instruction + 1 : transfer.receiver_instruction
                ]
                if transfer.receiver in programs
                else ()
            )
            can_group = not summarize_operations(sender_gap) and not summarize_operations(receiver_gap)

        if can_group:
            grouped[-1] = previous._replace(
                size=previous.size + transfer.size,
                sender_last_instruction=transfer.sender_last_instruction,
                receiver_last_instruction=transfer.receiver_last_instruction,
                count=previous.count + transfer.count,
            )
        else:
            grouped.append(transfer)
    return grouped


def collect_operation_notes(
    cores: list[int],
    transfers: list[Transfer],
    programs: dict[int, Program],
) -> dict[int, list[tuple[int, str]]]:
    notes: dict[int, list[tuple[int, str]]] = {}
    last_instructions = {core: -1 for core in cores}
    anchors = {core: -1 for core in cores}

    for transfer_index, transfer in enumerate(transfers):
        endpoints = {
            transfer.sender: (transfer.sender_instruction, transfer.sender_last_instruction),
            transfer.receiver: (transfer.receiver_instruction, transfer.receiver_last_instruction),
        }
        for core, (instruction, last_instruction) in endpoints.items():
            if core not in last_instructions or instruction is None:
                continue
            summary = summarize_operations(
                programs[core].operations[last_instructions[core] + 1 : instruction]
            )
            if summary:
                notes.setdefault(anchors[core], []).append((core, summary))
            last_instructions[core] = last_instruction
            anchors[core] = transfer_index

    for core in cores:
        summary = summarize_operations(programs[core].operations[last_instructions[core] + 1 :])
        if summary:
            notes.setdefault(anchors[core], []).append((core, summary))
    return notes


def note_height(summary: str) -> float:
    return 2.5 + 1.6 * (summary.count("\\n") + 1)


def parallel_note_order(
    notes: list[tuple[int, str]],
    cores: list[int],
) -> list[tuple[int, str]]:
    positions = {core: position for position, core in enumerate(cores)}
    packed: list[list[tuple[int, str]]] = []
    for note in sorted(notes, key=lambda item: positions[item[0]]):
        row = next(
            (
                row
                for row in packed
                if positions[note[0]] - positions[row[-1][0]] > 1
            ),
            None,
        )
        if row is None:
            row = []
            packed.append(row)
        row.append(note)
    return [note for row in packed for note in row]


def render_text(cores: list[int], transfers: list[Transfer], programs: dict[int, Program]) -> str:
    aliases = {core: f"C{core}" for core in cores}
    positions = {core: position for position, core in enumerate(cores)}
    lines = [f'participant "Core {core}" as {aliases[core]}' for core in cores]
    first_receives = {
        core: next((event.ordinal for event in programs[core].events if event.op == "recv"), None)
        for core in cores
    }
    active = {core for core in cores if not programs[core].starts_inactive}
    lines.extend(f"activate {aliases[core]}" for core in cores if core in active)
    notes = collect_operation_notes(cores, transfers, programs)
    availability = [0.0] * (2 * len(cores) + 1)
    cursor = 0.0

    def emit(line: str, left: int, right: int, height: float) -> None:
        nonlocal cursor
        start = max(availability[left : right + 1])
        gap = start - cursor
        if abs(gap) >= 0.05:
            value = f"{gap:.1f}".rstrip("0").rstrip(".")
            lines.append(f"space {value}")
        lines.append(line)
        cursor = start + height
        availability[left : right + 1] = [cursor] * (right - left + 1)

    def emit_notes(anchor: int) -> None:
        for core, summary in parallel_note_order(notes.get(anchor, []), cores):
            center = 2 * positions[core] + 1
            emit(f"note over {aliases[core]}:{summary}", center - 1, center + 1, note_height(summary))

    emit_notes(-1)

    for transfer_index, transfer in enumerate(transfers):
        label = f"{transfer.size} B"
        if transfer.count > 1:
            label = f"{transfer.count} sends, {label}"
        if transfer.sender in aliases and transfer.receiver in aliases:
            sender = 2 * positions[transfer.sender] + 1
            receiver = 2 * positions[transfer.receiver] + 1
            line = f"{aliases[transfer.sender]}->(1){aliases[transfer.receiver]}:{label}"
            left, right = sorted((sender, receiver))
        elif transfer.receiver in aliases:
            line = f"[->(1){aliases[transfer.receiver]}:{label}"
            left, right = 0, 2 * positions[transfer.receiver] + 1
        else:
            line = f"{aliases[transfer.sender]}->(1)]:{label}"
            left, right = 2 * positions[transfer.sender] + 1, len(availability) - 1
        emit(line, left, right, ARROW_HEIGHT)

        receiver = transfer.receiver
        if (
            receiver in aliases
            and receiver not in active
            and transfer.receiver_ordinal == first_receives[receiver]
        ):
            lines.append(f"activate {aliases[receiver]}")
            active.add(receiver)

        sender = transfer.sender
        if (
            sender in active
            and programs[sender].ends_with_send
            and transfer.sender_ordinal + transfer.count == len(programs[sender].events)
        ):
            lines.append(f"deactivate {aliases[sender]}")
            active.remove(sender)

        emit_notes(transfer_index)

    lines.extend(f"deactivateafter {aliases[core]}" for core in cores if core in active)
    return "\n".join(lines) + "\n"


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Generate SequenceDiagram.org text for PIM communication and intervening work."
    )
    parser.add_argument("pim_dir", type=Path, help="Directory containing core_<id>.json or core_<id>.pim files")
    selection = parser.add_mutually_exclusive_group(required=True)
    selection.add_argument("--cores", nargs="+", type=int, help="Core lifelines to display, in column order")
    selection.add_argument("--all-cores", action="store_true", help="Display every used core, ordered by core ID")
    parser.add_argument("--format", choices=("auto", "json", "pim"), default="auto")
    parser.add_argument("-o", "--output", type=Path, help="Text output path (default: stdout)")
    args = parser.parse_args()

    try:
        if not args.pim_dir.is_dir():
            raise ValueError(f"{args.pim_dir}: not a directory")
        if args.cores is not None and len(set(args.cores)) != len(args.cores):
            raise ValueError("--cores contains duplicates")
        programs = read_programs(args.pim_dir, args.format)
        cores = (
            [core for core in sorted(programs) if programs[core].operations]
            if args.all_cores
            else args.cores
        )
        missing = [core for core in cores if core not in programs]
        if missing:
            raise ValueError(f"missing artifact for selected core(s): {', '.join(map(str, missing))}")
        transfers = visible_transfers(match_transfers(programs), set(cores), programs)
        diagram = render_text(cores, transfers, programs)
        if args.output:
            args.output.write_text(diagram, encoding="utf-8")
        else:
            sys.stdout.write(diagram)
    except (OSError, ValueError) as error:
        parser.error(str(error))
    return 0


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