"""Command-line interface for training, inference, preencoding, and data inspection."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import sys
import tempfile
from collections.abc import Sequence
from pathlib import Path
from typing import Any

from solarwm import __version__
from solarwm.config import load_config
from solarwm.config.routes import supported_routes, validate_route
from solarwm.data.index import inventory, read_index
from solarwm.data.sampling import (
    CanonicalSampler,
    ReaderIdentity,
    SamplingConfig,
    plan_fingerprint,
)
from solarwm.errors import SolarWMError
from solarwm.runtime.create_only import write_file_create_only
from solarwm.runtime.images import probe_python_runtime
from solarwm.runtime.output_layout import (
    camera_inference_output_layout,
    invocation_output_dir,
)
from solarwm.runtime.provenance import build_launch_manifest, reject_inline_secrets


def _common_config(parser: argparse.ArgumentParser) -> None:
    parser.add_argument("--config", required=True, type=Path)
    parser.add_argument(
        "--set",
        action="append",
        default=[],
        metavar="KEY=VALUE",
        help="override one dotted config value; may be repeated",
    )


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="solarwm")
    parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
    commands = parser.add_subparsers(dest="command", required=True)

    for action in ("train", "infer", "preencode"):
        child = commands.add_parser(action, help=f"run the unified {action} entrypoint")
        _common_config(child)

    config = commands.add_parser("config", help="resolve and inspect a run config")
    config_commands = config.add_subparsers(dest="config_command", required=True)
    resolve = config_commands.add_parser("resolve")
    _common_config(resolve)
    resolve.add_argument("--output", type=Path)
    config_commands.add_parser("routes")

    data = commands.add_parser("data", help="inspect canonical index semantics")
    data_commands = data.add_subparsers(dest="data_command", required=True)
    inspect = data_commands.add_parser("inspect")
    inspect.add_argument("index", type=Path)
    plan = data_commands.add_parser("plan")
    plan.add_argument("index", type=Path)
    plan.add_argument("--seed", type=int, required=True)
    plan.add_argument("--pixel-frames", type=int, required=True)
    plan.add_argument("--rank", type=int, default=0)
    plan.add_argument("--world-size", type=int, default=1)
    plan.add_argument("--worker-id", type=int, default=0)
    plan.add_argument("--num-workers", type=int, default=1)
    plan.add_argument("--epoch", type=int, default=0)
    plan.add_argument("--limit", type=int, default=16)
    plan.add_argument("--shuffle-buffer", type=int, default=4096)

    environment = commands.add_parser("environment", help="inspect the current runtime")
    environment_commands = environment.add_subparsers(dest="environment_command", required=True)
    environment_commands.add_parser("probe")

    return parser


def _write_json(value: Any, output: Path | None = None) -> None:
    payload = json.dumps(value, indent=2, sort_keys=True, ensure_ascii=False) + "\n"
    if output is None:
        sys.stdout.write(payload)
        return
    output.parent.mkdir(parents=True, exist_ok=True)
    temporary: str | None = None
    try:
        with tempfile.NamedTemporaryFile(
            mode="w",
            encoding="utf-8",
            dir=output.parent,
            prefix=f".{output.name}.",
            suffix=".tmp",
            delete=False,
        ) as handle:
            temporary = handle.name
            handle.write(payload)
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(temporary, output)
        temporary = None
    finally:
        if temporary is not None:
            Path(temporary).unlink(missing_ok=True)


def _run_backend(args: argparse.Namespace) -> int:
    resolved = load_config(args.config, args.set)
    reject_inline_secrets(resolved.values)
    configured_action = str(resolved.values["action"])
    if configured_action != args.command:
        raise SolarWMError(
            f"command {args.command!r} does not match config action {configured_action!r}"
        )
    route = validate_route(resolved.values)
    runtime = resolved.values.get("runtime", {})

    from solarwm.backends import load_backend

    backend = load_backend(route.family)
    backend.validate_config(resolved.values)
    if not str(runtime.get("output_dir", "")).strip():
        raise SolarWMError("runtime.output_dir is required")
    publication_layout = camera_inference_output_layout(resolved.values)
    output_dir = invocation_output_dir(resolved.values)
    launcher_rank = int(os.environ.get("RANK", "0"))
    if launcher_rank == 0:
        if publication_layout is not None:
            publication_layout.publish_root.mkdir(parents=True, exist_ok=True)
            output_dir.parent.mkdir(parents=True, exist_ok=True)
            try:
                output_dir.mkdir()
            except FileExistsError as exc:
                raise SolarWMError(f"camera inference run already exists: {output_dir}") from exc
        else:
            output_dir.mkdir(parents=True, exist_ok=True)
            # A terminal result describes only the current invocation. Remove
            # a result left by an earlier legacy invocation before work starts.
            (output_dir / "run-result.json").unlink(missing_ok=True)
        resolved.write_json(output_dir / "resolved-config.json")
        manifest = build_launch_manifest(
            config=resolved.values,
            source_config=resolved.path,
            source_digest=resolved.source_digest,
            resolved_digest=resolved.resolved_digest,
            route=route.key,
            repository=Path(__file__).resolve().parents[2],
        )
        _write_json(manifest, output_dir / "launch-manifest.json")
    operation = getattr(backend, args.command)
    try:
        result = int(operation(resolved.values) or 0)
    except Exception as exc:
        if launcher_rank == 0:
            _write_json(
                {
                    "schema": "solarwm.run-result.v1",
                    "status": "failed",
                    "error_type": type(exc).__name__,
                    "message": str(exc),
                },
                output_dir / "run-result.json",
            )
        raise
    if launcher_rank == 0:
        _write_json(
            {"schema": "solarwm.run-result.v1", "status": "complete", "exit_code": result},
            output_dir / "run-result.json",
        )
        if publication_layout is not None and result == 0:
            publication_complete = output_dir / "publication" / "COMPLETE.json"
            if not publication_complete.is_file():
                raise SolarWMError(
                    "camera inference returned without a publication completion marker"
                )
            publication_payload = publication_complete.read_bytes()
            run_result_payload = (output_dir / "run-result.json").read_bytes()
            complete_payload = (
                json.dumps(
                    {
                        "schema": "solarwm.camera-inference-run-complete.v1",
                        "layout": publication_layout.layout,
                        "run_id": publication_layout.run_id,
                        "resolved_config_digest": resolved.resolved_digest,
                        "publication_complete_digest": hashlib.blake2s(
                            publication_payload
                        ).hexdigest(),
                        "run_result_digest": hashlib.blake2s(run_result_payload).hexdigest(),
                    },
                    allow_nan=False,
                    sort_keys=True,
                    separators=(",", ":"),
                    ensure_ascii=False,
                )
                + "\n"
            ).encode("utf-8")
            write_file_create_only(
                output_dir / "COMPLETE.json",
                complete_payload,
                error_type=SolarWMError,
                label="camera inference run completion marker",
            )
    return result


def _config_command(args: argparse.Namespace) -> int:
    if args.config_command == "routes":
        _write_json([route.key for route in supported_routes()])
        return 0
    resolved = load_config(args.config, args.set)
    reject_inline_secrets(resolved.values)
    route = validate_route(resolved.values)
    from solarwm.backends import load_backend

    load_backend(route.family).validate_config(resolved.values)
    payload = {
        "source": str(resolved.path),
        "source_digest": resolved.source_digest,
        "resolved_digest": resolved.resolved_digest,
        "route": route.key,
        "config": resolved.values,
    }
    _write_json(payload, args.output)
    return 0


def _data_command(args: argparse.Namespace) -> int:
    rows = read_index(args.index)
    report = inventory(args.index, rows)
    if args.data_command == "inspect":
        _write_json(
            {
                "index": str(args.index.resolve()),
                "rows": report.rows,
                "virtual_occurrences": report.virtual_occurrences,
                "ordered_sample_id_digest": report.ordered_sample_id_digest,
                "ordered_row_digest": report.ordered_row_digest,
                "decompressed_digest": report.decompressed_digest,
            }
        )
        return 0

    config = SamplingConfig(
        seed=args.seed,
        pixel_frames=args.pixel_frames,
        shuffle_buffer=args.shuffle_buffer,
    )
    identity = ReaderIdentity(
        rank=args.rank,
        world_size=args.world_size,
        worker_id=args.worker_id,
        num_workers=args.num_workers,
    )
    plans = list(CanonicalSampler(rows, config, identity).iter_epoch(args.epoch))
    selected = plans[: max(0, args.limit)]
    _write_json(
        {
            "full_plan_fingerprint": plan_fingerprint(plans),
            "full_plan_count": len(plans),
            "shown": [
                {
                    "sample_id": plan.sample_id,
                    "key": plan.key,
                    "shard": plan.shard,
                    "repeat_ordinal": plan.repeat_ordinal,
                    "start_frame": plan.start_frame,
                    "source_frame_indices": plan.source_frame_indices,
                }
                for plan in selected
            ],
        }
    )
    return 0


def _environment_command(args: argparse.Namespace) -> int:
    if args.environment_command == "probe":
        _write_json(probe_python_runtime())
        return 0
    raise SolarWMError(f"unhandled environment command {args.environment_command!r}")


def _run_backend_or_leave(args: argparse.Namespace) -> int:
    """Run a backend command; a rank of a distributed launch that fails exits outright.

    A rank that returns from a failure still joins its threads at interpreter shutdown, and one
    that never ends leaves it alive: the launcher never sees it fail, and every other
    rank waits out the collective timeout. Leaving at once hands the failure to the launcher.
    """

    if int(os.environ.get("WORLD_SIZE", "1")) <= 1:
        return _run_backend(args)
    try:
        return _run_backend(args)
    except (SolarWMError, OSError, ValueError) as exc:
        print(f"solarwm: error: {exc}", file=sys.stderr)
        status = 2
    except BaseException:
        import traceback

        traceback.print_exc()
        status = 1
    sys.stdout.flush()
    sys.stderr.flush()
    os._exit(status)


def main(argv: Sequence[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    try:
        if args.command in {"train", "infer", "preencode"}:
            return _run_backend_or_leave(args)
        if args.command == "config":
            return _config_command(args)
        if args.command == "data":
            return _data_command(args)
        if args.command == "environment":
            return _environment_command(args)
        raise SolarWMError(f"unhandled command {args.command!r}")
    except (SolarWMError, OSError, ValueError) as exc:
        print(f"solarwm: error: {exc}", file=sys.stderr)
        return 2


if __name__ == "__main__":  # pragma: no cover
    raise SystemExit(main())