"""Canonical generation matrix construction and TSV serialization.

The matrix is deliberately a small, dependency-light artifact.  It contains
only supported registry cells; unsupported cells are written to a separate
inventory so a skipped experiment is distinguishable from a missing row.
"""

from __future__ import annotations

import argparse
import csv
from dataclasses import asdict, dataclass
import os
from pathlib import Path
from typing import Iterable

from dlb.registry import ExperimentRegistry, load_registry, step_grid_for_model


MATRIX_SCHEMA = "dlb-generation-matrix-v1"
MATRIX_COLUMNS = (
    "task_id",
    "category",
    "model",
    "dataset",
    "steps",
    "sample_count",
    "seed",
    "environment",
    "adapter",
    "source",
    "provenance",
    "sample_dir",
    "metrics_path",
    "timing_path",
)
UNSUPPORTED_COLUMNS = ("status", "model", "dataset", "category", "reason")
SUPPORTED_STATUS = "supported"


@dataclass(frozen=True)
class MatrixTask:
    """One supported model/dataset/step generation request."""

    task_id: str
    category: str
    model: str
    dataset: str
    steps: int
    sample_count: int
    seed: int
    environment: str
    adapter: str
    source: str
    provenance: str
    sample_dir: str
    metrics_path: str
    timing_path: str

    @property
    def model_id(self) -> str:
        return self.model

    @property
    def dataset_id(self) -> str:
        return self.dataset

    @property
    def step_count(self) -> int:
        return self.steps

    def row(self) -> dict[str, str]:
        value = asdict(self)
        return {key: str(value[key]) for key in MATRIX_COLUMNS}


def _root_path(root: Path | None) -> Path:
    return (root or Path(".")).resolve()


def _task_paths(root: Path, model: str, dataset: str, steps: int) -> tuple[str, str, str]:
    sample_dir = root / "results" / "samples" / dataset / model / f"steps_{steps}"
    metrics = root / "results" / "metrics" / dataset / model / f"steps_{steps}" / "metrics.json"
    timing = root / "results" / "timing" / dataset / model / f"steps_{steps}" / "timing.json"
    return tuple(path.as_posix() for path in (sample_dir, metrics, timing))


def build_matrix(
    registry: ExperimentRegistry,
    *,
    root: Path | None = None,
    sample_count: int = 1024,
    seed: int = 42,
) -> list[MatrixTask]:
    """Expand every supported registry cell over its declared step grid."""

    if sample_count <= 0:
        raise ValueError("sample_count must be positive")
    root_path = _root_path(root)
    tasks: list[MatrixTask] = []
    for model_id, model in registry.models.items():
        for dataset_id, support in model.datasets.items():
            if support.status != SUPPORTED_STATUS:
                continue
            for steps in step_grid_for_model(registry, model_id):
                task_id = f"{model_id}-{dataset_id}-steps-{steps}"
                sample_dir, metrics_path, timing_path = _task_paths(
                    root_path, model_id, dataset_id, steps
                )
                tasks.append(
                    MatrixTask(
                        task_id=task_id,
                        category=model.category,
                        model=model_id,
                        dataset=dataset_id,
                        steps=steps,
                        sample_count=sample_count,
                        seed=seed,
                        environment=model.environment,
                        adapter=model.adapter,
                        source=model.source,
                        provenance=support.provenance or "",
                        sample_dir=sample_dir,
                        metrics_path=metrics_path,
                        timing_path=timing_path,
                    )
                )
    tasks.sort(key=lambda item: (item.model, item.dataset, item.steps))
    _validate_tasks(tasks)
    return tasks


def unsupported_inventory(registry: ExperimentRegistry) -> list[dict[str, str]]:
    """Return explicit unsupported cells without expanding them into steps."""

    records: list[dict[str, str]] = []
    for model_id, model in registry.models.items():
        for dataset_id, support in model.datasets.items():
            if support.status == SUPPORTED_STATUS:
                continue
            records.append(
                {
                    "status": "unsupported",
                    "model": model_id,
                    "dataset": dataset_id,
                    "category": model.category,
                    "reason": support.reason or "no reason recorded",
                }
            )
    records.sort(key=lambda item: (item["model"], item["dataset"]))
    return records


def write_unsupported_inventory(
    path: Path, records: Iterable[dict[str, str]], *, schema: str = MATRIX_SCHEMA
) -> Path:
    """Write the explicit unsupported model/dataset inventory as TSV."""

    path = Path(path).absolute()
    values = list(records)
    for record in values:
        if tuple(record) != UNSUPPORTED_COLUMNS:
            raise ValueError("unsupported inventory columns do not match the canonical schema")
        if any(
            "\t" in value or "\n" in value or "\r" in value
            for value in record.values()
        ):
            raise ValueError("unsupported inventory contains a control character")
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(f".{path.name}.{os.getpid()}.partial")
    try:
        with temporary.open("w", encoding="utf-8", newline="") as handle:
            handle.write(f"# schema={schema}-unsupported\n")
            writer = csv.DictWriter(
                handle, fieldnames=UNSUPPORTED_COLUMNS, delimiter="\t", lineterminator="\n"
            )
            writer.writeheader()
            writer.writerows(values)
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(temporary, path)
    except BaseException:
        temporary.unlink(missing_ok=True)
        raise
    return path


def _validate_tasks(tasks: Iterable[MatrixTask]) -> None:
    seen: set[str] = set()
    previous: tuple[str, str, int] | None = None
    for task in tasks:
        if task.task_id in seen:
            raise ValueError(f"duplicate matrix task ID: {task.task_id}")
        seen.add(task.task_id)
        key = (task.model, task.dataset, task.steps)
        if previous is not None and key < previous:
            raise ValueError("matrix tasks are not in stable order")
        previous = key
        if task.category not in {"many", "few", "fixed_1024"}:
            raise ValueError(f"invalid task category: {task.category}")
        if task.steps <= 0 or task.sample_count <= 0:
            raise ValueError("steps and sample_count must be positive")
        if any("\t" in value or "\n" in value or "\r" in value for value in task.row().values()):
            raise ValueError(f"matrix task contains a control character: {task.task_id}")


def write_matrix(path: Path, tasks: Iterable[MatrixTask]) -> Path:
    """Write a deterministic, versioned TSV matrix."""

    path = Path(path).absolute()
    task_list = list(tasks)
    _validate_tasks(task_list)
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(f".{path.name}.{os.getpid()}.partial")
    try:
        with temporary.open("w", encoding="utf-8", newline="") as handle:
            handle.write(f"# schema={MATRIX_SCHEMA}\n")
            writer = csv.DictWriter(
                handle, fieldnames=MATRIX_COLUMNS, delimiter="\t", lineterminator="\n"
            )
            writer.writeheader()
            for task in task_list:
                writer.writerow(task.row())
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(temporary, path)
    except BaseException:
        temporary.unlink(missing_ok=True)
        raise
    return path


def read_matrix(path: Path) -> list[MatrixTask]:
    """Read and validate a matrix TSV without evaluating any field as code."""

    path = Path(path)
    if path.is_symlink() or not path.is_file():
        raise ValueError(f"matrix TSV is missing or unsafe: {path}")
    with path.open("r", encoding="utf-8", newline="") as handle:
        schema = handle.readline().rstrip("\r\n")
        if schema != f"# schema={MATRIX_SCHEMA}":
            raise ValueError("matrix TSV has an unsupported schema header")
        reader = csv.DictReader(handle, delimiter="\t")
        if tuple(reader.fieldnames or ()) != MATRIX_COLUMNS:
            raise ValueError("matrix TSV columns do not match the canonical schema")
        tasks: list[MatrixTask] = []
        for row_number, row in enumerate(reader, start=3):
            if None in row or any(value is None for value in row.values()):
                raise ValueError(f"matrix TSV row {row_number} has missing fields")
            try:
                task = MatrixTask(
                    task_id=row["task_id"],
                    category=row["category"],
                    model=row["model"],
                    dataset=row["dataset"],
                    steps=int(row["steps"]),
                    sample_count=int(row["sample_count"]),
                    seed=int(row["seed"]),
                    environment=row["environment"],
                    adapter=row["adapter"],
                    source=row["source"],
                    provenance=row["provenance"],
                    sample_dir=row["sample_dir"],
                    metrics_path=row["metrics_path"],
                    timing_path=row["timing_path"],
                )
            except (KeyError, TypeError, ValueError) as error:
                raise ValueError(f"invalid matrix TSV row {row_number}") from error
            tasks.append(task)
    _validate_tasks(tasks)
    if len({task.task_id for task in tasks}) != len(tasks):
        raise ValueError("matrix TSV contains duplicate task IDs")
    return tasks


def validate_matrix(path: Path, registry: ExperimentRegistry | None = None) -> list[MatrixTask]:
    """Validate a matrix and optionally compare it with the canonical registry."""

    tasks = read_matrix(path)
    if registry is not None:
        expected = build_matrix(
            registry,
            root=Path("."),
            sample_count=tasks[0].sample_count if tasks else 1024,
            seed=tasks[0].seed if tasks else 42,
        )
        actual_keys = [(task.model, task.dataset, task.steps) for task in tasks]
        expected_keys = [(task.model, task.dataset, task.steps) for task in expected]
        if actual_keys != expected_keys:
            raise ValueError("matrix TSV coverage differs from the canonical registry")
    return tasks


def _main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--root", type=Path, default=Path.cwd())
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--sample-count", type=int, default=1024)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--unsupported-output", type=Path)
    parser.add_argument("--validate", action="store_true")
    args = parser.parse_args(argv)
    root = args.root.resolve()
    registry = load_registry(root / "configs" / "experiments.yaml")
    if args.validate:
        tasks = validate_matrix(args.output, registry)
    else:
        tasks = build_matrix(
            registry, root=root, sample_count=args.sample_count, seed=args.seed
        )
        write_matrix(args.output, tasks)
        write_unsupported_inventory(
            args.unsupported_output
            or args.output.with_name("unsupported.tsv"),
            unsupported_inventory(registry),
        )
    print(f"matrix={args.output} tasks={len(tasks)}")
    return 0


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