"""Generate versioned trial metrics and paper source-data tables."""

from __future__ import annotations

import argparse
import csv
import json
import os
import tempfile
from pathlib import Path
from typing import Any, Mapping, Sequence

import numpy as np

from experiments.hashing import file_sha256, stable_hash, to_jsonable
from experiments.io import atomic_write_json, atomic_write_jsonl, utc_now
from experiments.plan import load_document
from experiments.schema import SCHEMA_VERSION, STORAGE_FORMATS
from experiments.validate import require_valid_batch

from .metrics import derive_trial_metrics


def _method_selected(
    family: str,
    family_configuration: Mapping[str, Any],
    method_id: str,
) -> bool:
    """Apply explicit per-family method filters without silent typos."""
    include = family_configuration.get("include_methods")
    exclude = family_configuration.get("exclude_methods")
    for name, value in (
        ("include_methods", include),
        ("exclude_methods", exclude),
    ):
        if value is not None and (
            not isinstance(value, list)
            or any(not isinstance(item, str) or not item for item in value)
        ):
            raise ValueError(f"{family}.{name} must be a list of method IDs")
    if include is not None and method_id not in include:
        return False
    if exclude is not None and method_id in exclude:
        return False
    return True


def _trial_metric_configuration(
    configuration: Mapping[str, Any],
    method_id: str,
) -> dict[str, Any] | None:
    enabled = configuration.get("enabled")
    if not isinstance(enabled, list) or not enabled:
        raise ValueError("metric configuration needs a non-empty enabled list")
    selected = []
    trial_configuration: dict[str, Any] = {}
    for family in enabled:
        if not isinstance(family, str) or not family:
            raise ValueError("enabled metric family names must be non-empty strings")
        family_configuration = configuration.get(family, {})
        if not isinstance(family_configuration, Mapping):
            raise ValueError(f"{family} metric configuration must be a mapping")
        if _method_selected(family, family_configuration, method_id):
            selected.append(family)
            trial_configuration[family] = dict(family_configuration)
    if not selected:
        return None
    trial_configuration["enabled"] = selected
    return trial_configuration


def _scalar_cell(value: Any) -> Any:
    if value is None or isinstance(value, (str, int, float, bool)):
        return value
    return json.dumps(
        to_jsonable(value),
        sort_keys=True,
        separators=(",", ":"),
        ensure_ascii=False,
    )


def _identity_row(trial: Mapping[str, Any]) -> dict[str, Any]:
    trajectory = trial["trajectory"]
    method = trial["method"]
    row = {
        "trial_id": trial["trial_id"],
        "pair_id": trial["pair_id"],
        "study_id": trial["study_id"],
        "split": trial["split"],
        "method_id": method["method_id"],
        "trajectory_id": trajectory["trajectory_id"],
        "trajectory_family": trajectory.get("family", ""),
        "replicate": trial["replicate"],
    }
    for name, value in sorted(trial.get("factors", {}).items()):
        row[f"factor_{name}"] = _scalar_cell(value)
    return row


def _atomic_write_csv(path: Path, rows: Sequence[Mapping[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    identity_order = [
        "trial_id",
        "pair_id",
        "study_id",
        "split",
        "method_id",
        "trajectory_id",
        "trajectory_family",
        "replicate",
    ]
    all_fields = {key for row in rows for key in row}
    fieldnames = [name for name in identity_order if name in all_fields]
    fieldnames.extend(sorted(all_fields - set(fieldnames)))
    fd, temporary_name = tempfile.mkstemp(
        prefix=f".{path.name}.", suffix=".tmp", dir=path.parent
    )
    temporary = Path(temporary_name)
    try:
        with os.fdopen(fd, "w", encoding="utf-8", newline="") as stream:
            writer = csv.DictWriter(stream, fieldnames=fieldnames)
            writer.writeheader()
            for row in rows:
                writer.writerow(
                    {name: _scalar_cell(row.get(name)) for name in fieldnames}
                )
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temporary, path)
    except Exception:
        temporary.unlink(missing_ok=True)
        raise


def generate_paper_source_data(
    batch_dir: Path,
    metric_configuration: Mapping[str, Any],
    *,
    output_root: Path | None = None,
) -> dict[str, Any]:
    """Recompute all configured metrics and emit traceable source-data tables."""
    batch_dir = Path(batch_dir)
    require_valid_batch(batch_dir, require_complete=True)
    plan = load_document(batch_dir / "plan.json")
    output_root = batch_dir if output_root is None else Path(output_root)
    derived_dir = output_root / "derived"
    source_dir = output_root / "paper" / "source_data"
    derived_dir.mkdir(parents=True, exist_ok=True)
    source_dir.mkdir(parents=True, exist_ok=True)

    rows: list[dict[str, Any]] = []
    input_files: dict[str, str] = {}
    for trial in plan["trials"]:
        trial_dir = batch_dir / "raw" / trial["trial_id"]
        sample_path = trial_dir / "samples.npz"
        input_files[str(sample_path.relative_to(batch_dir))] = file_sha256(sample_path)
        with np.load(sample_path, allow_pickle=False) as archive:
            samples = {name: archive[name] for name in archive.files}
        method_id = trial["method"]["method_id"]
        trial_configuration = _trial_metric_configuration(
            metric_configuration,
            method_id,
        )
        metrics = (
            {}
            if trial_configuration is None
            else derive_trial_metrics(samples, trial_configuration)
        )
        rows.append({**_identity_row(trial), **metrics})

    metric_path = derived_dir / "trial_metrics.jsonl"
    atomic_write_jsonl(metric_path, rows)
    table_files: dict[str, str] = {}
    family_row_counts: dict[str, int] = {}
    for family in metric_configuration["enabled"]:
        family_rows: list[dict[str, Any]] = []
        prefix = f"{family}_"
        for row in rows:
            if not any(key.startswith(prefix) for key in row):
                continue
            selected = {
                key: value
                for key, value in row.items()
                if key in {
                    "trial_id",
                    "pair_id",
                    "study_id",
                    "split",
                    "method_id",
                    "trajectory_id",
                    "trajectory_family",
                    "replicate",
                }
                or key.startswith("factor_")
                or key.startswith(prefix)
            }
            family_rows.append(selected)
        if not family_rows:
            raise ValueError(
                f"metric family {family!r} selected no methods in this plan"
            )
        table_path = source_dir / f"{family}.csv"
        _atomic_write_csv(table_path, family_rows)
        table_files[str(table_path.relative_to(output_root))] = file_sha256(table_path)
        family_row_counts[family] = len(family_rows)

    manifest = {
        "kind": "paper_artifact_manifest",
        "schema_version": SCHEMA_VERSION,
        "created_utc": utc_now(),
        "source_batch": str(batch_dir),
        "source_plan_hash": plan["plan_hash"],
        "metric_configuration": to_jsonable(metric_configuration),
        "metric_configuration_hash": stable_hash(
            metric_configuration, prefix="metric-configuration"
        ),
        "row_count": len(rows),
        "family_row_counts": family_row_counts,
        "row_hash": stable_hash(rows, prefix="trial-metric-rows"),
        "storage_formats": {
            "trial_metrics": STORAGE_FORMATS["trial_metrics"],
            "paper_source_data": STORAGE_FORMATS["paper_source_data"],
        },
        "inputs": input_files,
        "outputs": {
            str(metric_path.relative_to(output_root)): file_sha256(metric_path),
            **table_files,
        },
    }
    manifest_path = output_root / "paper" / "artifact_manifest.json"
    atomic_write_json(manifest_path, manifest)
    return manifest


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Recompute trial metrics and paper source-data tables"
    )
    parser.add_argument("--batch-dir", type=Path, required=True)
    parser.add_argument("--metric-config", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, default=None)
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    configuration = load_document(args.metric_config)
    manifest = generate_paper_source_data(
        args.batch_dir,
        configuration,
        output_root=args.output_root,
    )
    print(json.dumps(manifest, indent=2, ensure_ascii=False, sort_keys=True))
    return 0


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