"""Adapter for the pinned official RDLM LM1B SDE sampler."""

from pathlib import Path
import sys
from typing import Iterable

from dlb.adapters.base import AdapterError, BaseTeacherAdapter, CheckpointSelection
from dlb.runner import RunRequest
from dlb.schema import SampleRecord


class RDLMAdapter(BaseTeacherAdapter):
    identity = "dlb.adapters.rdlm:v1"
    upstream = "rdlm"
    supported_models = frozenset({"rdlm"})
    teacher_families = {"rdlm": "continuous_rdlm"}
    batch_sizes = {("rdlm", "lm1b"): 8}
    _LM1B_FILES = (
        "LM1B/checkpoint.pth",
        "LM1B/config.yaml",
        "LM1B/sde.pkl",
    )

    def _asset_paths(
        self, checkpoint: CheckpointSelection
    ) -> tuple[Path, Path, Path]:
        missing = [name for name in self._LM1B_FILES if name not in checkpoint.required_files]
        if missing:
            raise AdapterError(
                "canonical RDLM resource does not bind the saved LM1B trio: "
                + ", ".join(missing)
            )
        asset_root = checkpoint.path / "LM1B"
        return (
            asset_root / "checkpoint.pth",
            asset_root / "config.yaml",
            asset_root / "sde.pkl",
        )

    def render_command(
        self, request: RunRequest, run_dir: Path, *, dry_run: bool
    ) -> list[str]:
        root, length, _ = self._validate_request(request, run_dir)
        if length != 128:
            raise AdapterError(f"RDLM LM1B requires sequence length 128, found {length}")
        checkpoint = self._resolve_checkpoint(root, request, dry_run=dry_run)
        checkpoint_path, saved_config, saved_sde = self._asset_paths(checkpoint)
        entrypoint = root / "upstreams" / "rdlm" / "main.py"
        if not entrypoint.is_file() or entrypoint.is_symlink():
            raise AdapterError(f"pinned upstream entrypoint is missing or unsafe: {entrypoint}")
        capture_path = run_dir.resolve() / "upstream_token_ids.json"
        hydra_dir = run_dir.resolve() / "upstream_hydra"
        tokenizer_name = self._load_data_contract(root, request.dataset_id)["tokenizer"]
        tokenizer_revision = self._tokenizer_revision(root, tokenizer_name)
        tokenizer_snapshot = (
            root
            / "data"
            / "raw"
            / "huggingface"
            / "hub"
            / f"models--{tokenizer_name}"
            / "snapshots"
            / tokenizer_revision
        )
        arguments = [
            sys.executable,
            "-B",
            "-u",
            "-m",
            "dlb.adapters.capture",
            f"--upstream-entrypoint={entrypoint}",
            f"--capture-path={capture_path}",
            "--capture-kind=rdlm",
            f"--saved-config-path={saved_config}",
            f"--saved-sde-path={saved_sde}",
            f"--expected-samples={request.sample_count}",
            f"--data-config-path={root / 'artifacts' / 'data.yaml'}",
            f"--downloads-manifest-path={root / 'data' / 'manifests' / 'downloads.json'}",
            f"--dataset-id={request.dataset_id}",
            f"--tokenizer-snapshot={tokenizer_snapshot}",
            *self._conditional_capture_flags(request),
            "--",
            "run_mode=sample",
            "server=sample",
            "exp=sample_lm1b",
            "ngpus=1",
            "use_wandb=False",
            f"model_path={checkpoint_path}",
            f"seed={request.seed}",
            "sampling.predictor=grw",
            f"sampling.steps={request.step_count}",
            f"sampling.batch_per_gpu={min(8, request.sample_count)}",
            "eval.entropy=False",
            "eval.nll=False",
            "eval.gen_ppl=False",
            f"hydra.run.dir={hydra_dir}",
        ]
        self._validate_argv(arguments)
        return arguments

    def convert_outputs(
        self, request: RunRequest, run_dir: Path
    ) -> Iterable[SampleRecord]:
        return self._convert_capture_outputs(
            request, run_dir, retokenize_inexact_rows=True
        )


adapter = RDLMAdapter()