#!/usr/bin/env python3
"""Validated merge of worker claims into the tracked registries.

Phase 8 runs matching in parallel: two worker sessions match functions and stage
their results as claim files, and one coordinator merges them. Hand-merging would
put the whole trust model on the coordinator's care; this tool makes the merge a
*checked* operation instead.

## What it validates

Every claim is a row in the registry format `start<TAB>end<TAB>source[<TAB>overrides]`.
A claim is accepted only if all of these hold:

  * the row parses, and `start < end`, both inside the declared payload;
  * the **extent exists in the derived extents table and is graded `exact`**, with
    exactly the claimed end -- so a claim cannot smuggle in a hand-chosen end
    address (see `docs/PHASE7_EXTENTS.md`);
  * the source is a repo-relative `src/func_XXXXXXXX.c` path that exists;
  * the region does not overlap any region already in the registry, nor any other
    claim in this merge (including one from a different worker's claim file).

Symbol rows (`NAME<TAB>address[<TAB>gp]`) are merged too: a name that already
exists with a *different* address is rejected, and a name that already exists with
the same address is a no-op.

## What it does not do

It does not compile anything. Extent exactness, overlap and payload bounds are
*structural* evidence; the proof that a claim's bytes are right is the full-binary
gate. The intended workflow is therefore:

    sf3_merge apply ... --out-regions .run/p8/regions.candidate.tsv \
                        --out-symbols .run/p8/symbols.candidate.tsv
    tools/sf3_match gate --regions .run/p8/regions.candidate.tsv \
                         --symbols .run/p8/symbols.candidate.tsv ...
    # only on MATCH: promote the candidate over the tracked registry

so the tracked registry never contains an unverified claim.

## Writing

Nothing is written unless **every** claim validates: rejections are reported with
reasons and the exit status is 1. Accepted output is written atomically (temp file
plus rename) to the explicit `--out-regions` / `--out-symbols` paths, with the
input file's header comments preserved and data rows sorted by address. Destinations
are explicit rather than defaulted, because a merge is an in-place edit by
definition and the project does not overwrite blind.

Exit codes: 0 success, 1 rejected claims, 2 usage or environment error.
"""

from __future__ import annotations

import argparse
import os
from pathlib import Path
import re
import struct
import sys
from typing import Sequence


EXE_MAGIC = b"PS-X EXE"
HEADER_SIZE = 0x800
PAYLOAD_LMA = 0x800

REQUIRED_GRADE = "exact"
SOURCE_PATTERN = re.compile(r"^src/func_[0-9A-Fa-f]{8}\.c$")
NAME_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_.]*$")


class ToolError(Exception):
    """A usage or environment problem; maps to exit code 2."""


class Rejection(Exception):
    """A claim that cannot be merged; collected and reported, never written."""


def parse_hex(text: str, label: str) -> int:
    text = text.strip()
    if text.lower() in ("start", "address", "addr"):
        raise Rejection(
            f"{label}: looks like a column HEADER, not a claim row. Header lines in a "
            f"claims file must start with '#'."
        )
    return _parse_hex(text, label)


def _parse_hex(text: str, label: str) -> int:
    try:
        return int(text, 16)
    except ValueError as exc:
        raise ToolError(f"{label}: not a hex address: {text!r}") from exc


def require_file(path: Path, label: str) -> Path:
    if not path.is_file():
        raise ToolError(f"{label} is not a regular file: {path}")
    return path


def parse_psx_exe(header: bytes) -> tuple[int, int, int]:
    if len(header) < HEADER_SIZE:
        raise ToolError("executable is smaller than a PS-X EXE header")
    if header[:8] != EXE_MAGIC:
        raise ToolError("executable does not carry the PS-X EXE magic")
    entry, _gp, text_address, text_size = struct.unpack_from("<IIII", header, 0x10)
    if text_size == 0:
        raise ToolError("PS-X EXE header declares an empty payload")
    return entry, text_address, text_size


def split_rows(text: str) -> tuple[list[str], list[tuple[int, str]]]:
    """Return the header comment lines and `(line number, data line)` rows."""
    header: list[str] = []
    rows: list[tuple[int, str]] = []
    for number, raw in enumerate(text.splitlines(), 1):
        line = raw.strip()
        if not line:
            continue
        if line.startswith("#"):
            if not rows:
                header.append(line)
            continue
        rows.append((number, line))
    return header, rows


def load_extents(path: Path) -> dict[int, tuple[int, str]]:
    """Map each start with an extent to `(end, grade)`."""
    _header, rows = split_rows(path.read_text(encoding="utf-8"))
    extents: dict[int, tuple[int, str]] = {}
    for number, line in rows:
        fields = line.split("\t")
        if len(fields) != 7:
            raise ToolError(f"extents line {number}: expected seven fields")
        if fields[1] == "-":
            continue
        start = parse_hex(fields[0], f"extents line {number}")
        end = parse_hex(fields[1], f"extents line {number}")
        extents[start] = (end, fields[5])
    if not extents:
        raise ToolError("extents table contains no rows with an extent")
    return extents


def load_regions(path: Path) -> tuple[list[str], list[tuple[int, int, str, str]]]:
    header, rows = split_rows(path.read_text(encoding="utf-8"))
    regions: list[tuple[int, int, str, str]] = []
    for number, line in rows:
        fields = line.split("\t")
        if len(fields) not in (3, 4):
            raise ToolError(f"region line {number}: expected three or four fields")
        start = parse_hex(fields[0], f"region line {number}")
        end = parse_hex(fields[1], f"region line {number}")
        if end <= start:
            raise ToolError(f"region line {number}: end is not after start")
        regions.append((start, end, fields[2], fields[3] if len(fields) == 4 else ""))
    return header, regions


def load_symbols(path: Path) -> tuple[list[str], list[tuple[str, int, str]]]:
    header, rows = split_rows(path.read_text(encoding="utf-8"))
    symbols: list[tuple[str, int, str]] = []
    for number, line in rows:
        fields = line.split("\t")
        if len(fields) not in (2, 3):
            raise ToolError(f"symbol line {number}: expected two or three fields")
        name = fields[0].strip()
        if not NAME_PATTERN.match(name):
            raise ToolError(f"symbol line {number}: invalid symbol name {name!r}")
        address = parse_hex(fields[1], f"symbol line {number}")
        marker = fields[2].strip() if len(fields) == 3 else ""
        if marker not in ("", "gp"):
            raise ToolError(f"symbol line {number}: unknown marker {marker!r}")
        symbols.append((name, address, marker))
    return header, symbols


class Claim:
    """One claimed region, with the file and line it came from."""

    __slots__ = ("start", "end", "source", "overrides", "origin")

    def __init__(self, start: int, end: int, source: str, overrides: str, origin: str) -> None:
        self.start = start
        self.end = end
        self.source = source
        self.overrides = overrides
        self.origin = origin

    def describe(self) -> str:
        return f"{self.origin}: 0x{self.start:08X}..0x{self.end:08X} {self.source}"


# The per-region override keys the harness understands. Validated HERE, at merge
# time, so an invalid claim row is rejected before it can reach a candidate: a row
# carrying an unknown key used to merge silently and only fail later at the gate
# (Phase 11 incident: a worker put `md5=...` in the 4th column, which merged and then
# made `sf3_match gate` abort with "unknown override key 'md5'"). Fail fast.
_REGION_OPTION_KEYS = ("cc1", "as", "gp", "maspsx", "cc1bin")
# `cc1bin=NAME` selects an alternative cc1 BUILD for a region (Phase 12, developer-approved).
# It is a BUILD NAME, not a flag list -- a bare directory name under tools/old-gcc/ -- and it
# is RESTRICTED: `sf3_match` refuses a region whose original body has fewer than two
# function-exit jumps, so the lever can only be applied to the class the default cc1 provably
# cannot emit. Validated here too, so a claim row carrying it is checked at merge time.


def validate_overrides(text: str, origin: str) -> str:
    """Reject a claim row whose override field carries an unknown key.

    A lone `-` means "no overrides" -- the same absent-value convention the other
    tracked tables use -- and is normalised to the empty string.
    """
    if text.strip() == "-":
        return ""
    for token in text.split():
        if "=" not in token:
            raise Rejection(f"{origin}: expected 'key=value' override, got {token!r}")
        key = token.partition("=")[0]
        if key not in _REGION_OPTION_KEYS:
            raise Rejection(
                f"{origin}: unknown override key {key!r}; valid keys are "
                f"{', '.join(_REGION_OPTION_KEYS)}. Per-claim metadata such as a source "
                f"md5 belongs in report.tsv, not in the registry row."
            )
    return text


def load_claims(path: Path) -> list[Claim]:
    _header, rows = split_rows(path.read_text(encoding="utf-8"))
    claims: list[Claim] = []
    for number, line in rows:
        fields = line.split("\t")
        origin = f"{path}:{number}"
        if len(fields) not in (3, 4):
            raise Rejection(f"{origin}: expected three or four fields")
        start = parse_hex(fields[0], origin)
        end = parse_hex(fields[1], origin)
        overrides = validate_overrides(fields[3], origin) if len(fields) == 4 else ""
        claims.append(Claim(start, end, fields[2], overrides, origin))
    return claims


def merge(regions: Sequence[tuple[int, int, str, str]], claims: Sequence[Claim],
          extents: dict[int, tuple[int, str]], text_address: int, payload_end: int,
          root: Path, skip_registered: bool = False) -> tuple[list[tuple[int, int, str, str]], list[str], int]:
    """Validate every claim; return the merged regions, the rejections, and skipped count."""
    merged = list(regions)
    rejections: list[str] = []
    skipped = 0
    occupied = sorted((start, end) for start, end, _s, _o in regions)
    exact_rows = {(start, end, source) for start, end, source, _o in regions}

    def overlaps(start: int, end: int) -> tuple[int, int] | None:
        for other_start, other_end in occupied:
            if start < other_end and other_start < end:
                return other_start, other_end
        return None

    for claim in sorted(claims, key=lambda c: (c.start, c.end)):
        if skip_registered and (claim.start, claim.end, claim.source) in exact_rows:
            # A worker's claims file is cumulative, so a row that is already
            # registered exactly as claimed is a no-op, not a conflict.
            skipped += 1
            continue
        try:
            if claim.start < text_address or claim.end > payload_end:
                raise Rejection("outside the declared payload")
            derived = extents.get(claim.start)
            if derived is None:
                raise Rejection("no derived extent for this start address")
            if derived[0] != claim.end:
                raise Rejection(
                    f"claimed end 0x{claim.end:08X} but the derived extent is 0x{derived[0]:08X}")
            if derived[1] != REQUIRED_GRADE:
                raise Rejection(
                    f"extent is graded {derived[1]!r}, not {REQUIRED_GRADE!r}")
            if not SOURCE_PATTERN.match(claim.source):
                raise Rejection(f"source {claim.source!r} is not src/func_XXXXXXXX.c")
            source_path = (root / claim.source).resolve()
            if root.resolve() not in source_path.parents:
                raise Rejection(f"source {claim.source!r} escapes the repository root")
            if not source_path.is_file():
                raise Rejection(f"source file does not exist: {claim.source}")
            clash = overlaps(claim.start, claim.end)
            if clash is not None:
                raise Rejection(
                    f"overlaps 0x{clash[0]:08X}..0x{clash[1]:08X}")
        except Rejection as rejection:
            rejections.append(f"{claim.describe()}: {rejection}")
            continue
        merged.append((claim.start, claim.end, claim.source, claim.overrides))
        occupied.append((claim.start, claim.end))
        occupied.sort()

    merged.sort(key=lambda row: (row[0], row[1]))
    return merged, rejections, skipped


def merge_symbols(existing: Sequence[tuple[str, int, str]],
                  incoming: Sequence[tuple[str, int, str]]) -> tuple[list[tuple[str, int, str]], list[str]]:
    merged = list(existing)
    rejections: list[str] = []
    known = {name: address for name, address, _marker in existing}
    for name, address, marker in incoming:
        if name in known:
            if known[name] != address:
                rejections.append(
                    f"symbol {name}: already defined at 0x{known[name]:08X}, claim says 0x{address:08X}")
            continue
        known[name] = address
        merged.append((name, address, marker))
    merged.sort(key=lambda row: row[0])
    return merged, rejections


def format_rows(header: Sequence[str], rows: Sequence[str]) -> str:
    return "\n".join([*header, *rows]) + "\n"


def write_atomic(path: Path, text: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(path.name + ".tmp")
    temporary.write_text(text, encoding="ascii")
    os.replace(temporary, path)


def command_apply(args: argparse.Namespace) -> int:
    exe_path = require_file(args.exe, "executable")
    extents_path = require_file(args.extents, "function extents")
    regions_path = require_file(args.regions, "region registry")
    claim_paths = [require_file(path, "claim file") for path in args.claims]
    if not claim_paths:
        raise ToolError("at least one --claims file is required")
    root = Path(args.root).resolve()
    if not root.is_dir():
        raise ToolError(f"--root is not a directory: {root}")

    data = exe_path.read_bytes()
    _entry, text_address, text_size = parse_psx_exe(data)
    payload_end = text_address + text_size
    extents = load_extents(extents_path)
    region_header, regions = load_regions(regions_path)

    claims: list[Claim] = []
    rejections: list[str] = []
    for path in claim_paths:
        try:
            claims.extend(load_claims(path))
        except Rejection as rejection:
            rejections.append(str(rejection))

    merged_regions, claim_rejections, skipped = merge(regions, claims, extents, text_address,
                                                      payload_end, root, args.skip_registered)
    rejections.extend(claim_rejections)

    symbol_header: list[str] = []
    merged_symbols: list[tuple[str, int, str]] = []
    added_symbols = 0
    if args.symbols is not None:
        symbol_path = require_file(args.symbols, "symbol registry")
        symbol_header, existing_symbols = load_symbols(symbol_path)
        incoming: list[tuple[str, int, str]] = []
        for path in args.claims_symbols:
            _header, rows = split_rows(require_file(path, "claim symbol file").read_text(encoding="utf-8"))
            for number, line in rows:
                fields = line.split("\t")
                if len(fields) not in (2, 3):
                    rejections.append(f"{path}:{number}: expected two or three fields")
                    continue
                name = fields[0].strip()
                if not NAME_PATTERN.match(name):
                    rejections.append(f"{path}:{number}: invalid symbol name {name!r}")
                    continue
                incoming.append((name, parse_hex(fields[1], f"{path}:{number}"),
                                 fields[2].strip() if len(fields) == 3 else ""))
        before = len(existing_symbols)
        merged_symbols, symbol_rejections = merge_symbols(existing_symbols, incoming)
        rejections.extend(symbol_rejections)
        added_symbols = len(merged_symbols) - before
    if args.out_symbols is not None and args.symbols is None:
        raise ToolError("--out-symbols requires --symbols")

    added_regions = len(merged_regions) - len(regions)
    for rejection in rejections:
        print(f"reject: {rejection}", file=sys.stderr)

    print(f"claims={len(claims)}")
    print(f"accepted={len(claims) - len(claim_rejections) - skipped}")
    print(f"skipped_registered={skipped}")
    print(f"rejected={len(rejections)}")
    print(f"added_regions={added_regions}")
    print(f"added_symbols={added_symbols}")
    print(f"regions_total={len(merged_regions)}")
    if args.symbols is not None:
        print(f"symbols_total={len(merged_symbols)}")

    if rejections:
        print("result=REJECTED")
        return 1
    if args.dry_run:
        print("result=OK (dry run, nothing written)")
        return 0

    region_rows = [
        "\t".join((f"0x{start:08X}", f"0x{end:08X}", source, overrides)).rstrip("\t")
        for start, end, source, overrides in merged_regions
    ]
    write_atomic(args.out_regions, format_rows(region_header, region_rows))
    print(f"out_regions={args.out_regions}")
    if args.symbols is not None and args.out_symbols is not None:
        symbol_rows = [
            "\t".join((name, f"0x{address:08X}", marker)).rstrip("\t")
            for name, address, marker in merged_symbols
        ]
        write_atomic(args.out_symbols, format_rows(symbol_header, symbol_rows))
        print(f"out_symbols={args.out_symbols}")
    print("result=MERGED")
    return 0


def command_check_claims(args: argparse.Namespace) -> int:
    """Validate claim files WITHOUT merging, so a worker can self-check before reporting.

    WHY THIS EXISTS -- Phase 12. The coordinator's charter told every worker to stage
    `claims.tsv` as `range<TAB>source<TAB>md5<TAB>differing_bytes<TAB>result<TAB>options`
    with the range written `0xSTART..0xEND`. That format is not what this tool reads, so
    the FIRST merge of the phase was rejected with `expected three or four fields` -- and
    all four workers were already writing files in the wrong format.

    The tool was right and the prose was wrong. That is cookbook 179 exactly: a rule every
    worker must follow belongs in a TRACKED tool, not in a file each worker copies -- and a
    charter is a file each worker copies. The durable fix is not to correct the prose but to
    give the worker a command that answers "is my staging acceptable?" before the merge is
    attempted, so the next coordinator's prose error costs one command instead of a batch.

    Checks the format, the absence of duplicate starts, and that every claimed source file
    exists. It does NOT check extents, overlaps with the registry, or the whole-binary gate:
    those belong to `apply` and to `sf3_match gate`, and `--regions` here is only used to
    warn about a row that is already registered.
    """
    claim_paths = [require_file(path, "claim file") for path in args.claims]
    if not claim_paths:
        raise ToolError("at least one --claims file is required")
    root = Path(args.root).resolve()
    if not root.is_dir():
        raise ToolError(f"--root is not a directory: {root}")

    registered: list[tuple[int, int]] = []
    if args.regions is not None:
        registered = [(start, end) for start, end, _s, _o in load_regions(args.regions)[1]]

    total = 0
    problems = 0
    already = 0
    for path in claim_paths:
        try:
            claims = load_claims(path)
        except Rejection as exc:
            print(f"rejected: {exc}", file=sys.stderr)
            return 2
        print(f"{path}: {len(claims)} claim(s)")
        seen: dict[int, str] = {}
        for claim in claims:
            issues = []
            registered_note = ""
            source_path = root / claim.source
            if not source_path.is_file():
                issues.append("SOURCE MISSING")
            if claim.start in seen:
                issues.append(f"DUPLICATE START with {seen[claim.start]}")
            if any(start <= claim.start < end for start, end in registered):
                already += 1
                # INFORMATION, not a problem: `apply --skip-registered` treats an
                # already-registered row as a deliberate no-op, and a worker's claims
                # file is CUMULATIVE, so after the first merge it always contains some.
                # Reporting that as a problem makes `result=PROBLEMS`/exit 1 the normal
                # case for a healthy file, which teaches workers to ignore the exit code
                # -- and then the checker catches nothing at all. Worker B reported exactly
                # this after its first merge.
                registered_note = "  [already registered: skipped by apply]"
            seen[claim.start] = claim.origin
            problems += len(issues)
            flag = "ok" if not issues else "  <-- " + "; ".join(issues)
            print(f"  {claim.describe()}  overrides={claim.overrides or '-'}  {flag}{registered_note}")
        total += len(claims)

    print(f"claims={total} problems={problems} already_registered={already}")
    print("result=" + ("OK" if problems == 0 else "PROBLEMS"))
    return 0 if problems == 0 else 1


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    apply_parser = subparsers.add_parser("apply", help="validate and merge claim files")
    apply_parser.add_argument("--exe", required=True, type=Path)
    apply_parser.add_argument("--extents", required=True, type=Path)
    apply_parser.add_argument("--regions", required=True, type=Path)
    apply_parser.add_argument("--claims", required=True, type=Path, action="append", default=[])
    apply_parser.add_argument("--symbols", type=Path, default=None)
    apply_parser.add_argument("--claims-symbols", type=Path, action="append", default=[],
                              dest="claims_symbols")
    apply_parser.add_argument("--out-regions", required=True, type=Path)
    apply_parser.add_argument("--out-symbols", type=Path, default=None)
    apply_parser.add_argument("--root", type=Path, default=Path("."),
                              help="root the claim source paths are resolved against")
    apply_parser.add_argument("--dry-run", action="store_true", dest="dry_run",
                              help="validate only; write nothing")
    apply_parser.add_argument("--skip-registered", action="store_true", dest="skip_registered",
                              help="skip a claim that is already registered exactly as claimed")
    apply_parser.set_defaults(handler=command_apply)

    check_parser = subparsers.add_parser(
        "check-claims", help="validate claim files without merging (worker self-check)"
    )
    check_parser.add_argument("--claims", required=True, type=Path, action="append", default=[])
    check_parser.add_argument("--root", type=Path, default=Path("."),
                              help="root the claimed source paths are relative to")
    check_parser.add_argument("--regions", type=Path, default=None,
                              help="optional: warn about rows already in this registry")
    check_parser.set_defaults(handler=command_check_claims)

    return parser


def main(argv: Sequence[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    try:
        return args.handler(args)
    except ToolError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 2
    except OSError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 2


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