from __future__ import annotations

from dataclasses import replace
from pathlib import Path

import pytest

from dlb.adapters.candi import CANDIAdapter
from dlb.adapters.di4c import Di4CAdapter
from dlb.adapters.duo import DuoAdapter
from dlb.adapters.flm import FLMAdapter
from dlb.adapters.langflow import LangFlowAdapter
from dlb.adapters.mdlm import MDLMAdapter
from dlb.adapters.rdlm import RDLMAdapter
from dlb.adapters.sdtt import SDTTAdapter
from dlb.runner import RunRequest


ROOT = Path(__file__).parents[1]
MANIFEST = ROOT / "data/manifests/conditional-lm1b-c64.json"


def conditional_request(model: str, dataset: str = "lm1b", steps: int = 8) -> RunRequest:
    step_count = 1000 if model == "rdlm" else steps
    return RunRequest(
        run_id=f"{model}-{dataset}-steps-{step_count}",
        model_id=model,
        dataset_id=dataset,
        step_count=step_count,
        seed=42,
        sample_count=2048,
        generation_mode="conditional_prefix",
        conditioning_manifest=str(MANIFEST),
        conditioning_manifest_sha256="a" * 64,
        conditioning_config_sha256="b" * 64,
        prefix_length=64,
        evaluation_continuation_length=64,
        prompt_count=1024,
        diversity_prompt_count=256,
        completions_per_diversity_prompt=5,
        completion_schedule="c0:p0-1023;c1-4:p0-255",
        results_root=str(ROOT / "results/conditional"),
    )


def run_dir(item: RunRequest) -> Path:
    return (
        ROOT
        / "results"
        / "conditional"
        / "samples"
        / item.dataset_id
        / item.model_id
        / f"steps_{item.step_count}"
    )


def option(command: list[str], key: str) -> str:
    index = command.index(key)
    return command[index + 1]


def wrapper_option(command: list[str], key: str) -> str:
    prefix = key + "="
    values = [argument[len(prefix) :] for argument in command if argument.startswith(prefix)]
    assert len(values) == 1, (key, command)
    return values[0]


@pytest.mark.parametrize(
    ("model", "adapter"),
    [
        ("flm", FLMAdapter()),
        ("fmlm", FLMAdapter()),
        ("duo", DuoAdapter()),
        ("duo_dcd", DuoAdapter()),
        ("mdlm", MDLMAdapter()),
        ("candi", CANDIAdapter()),
        ("langflow", LangFlowAdapter()),
        ("rdlm", RDLMAdapter()),
    ],
)
def test_teacher_command_serializes_complete_conditioning_contract(model: str, adapter) -> None:
    request = conditional_request(model)

    command = adapter.render_command(request, run_dir(request), dry_run=True)

    assert wrapper_option(command, "--generation-mode") == "conditional_prefix"
    assert wrapper_option(command, "--conditioning-manifest") == str(MANIFEST)
    assert wrapper_option(command, "--conditioning-manifest-sha256") == "a" * 64
    assert wrapper_option(command, "--prefix-length") == "64"
    assert wrapper_option(command, "--prompt-count") == "1024"
    assert wrapper_option(command, "--diversity-prompt-count") == "256"
    assert wrapper_option(command, "--completions-per-diversity-prompt") == "5"
    assert wrapper_option(command, "--completion-schedule") == "c0:p0-1023;c1-4:p0-255"


@pytest.mark.parametrize(
    ("model", "adapter"),
    [
        ("mdlm_sdtt", SDTTAdapter()),
        ("mdlm_di4c", Di4CAdapter("mdlm")),
        ("duo_di4c", Di4CAdapter("duo")),
    ],
)
def test_distilled_command_serializes_complete_conditioning_contract(model: str, adapter) -> None:
    request = conditional_request(model)

    command = adapter.render_command(request, run_dir(request), dry_run=True)

    assert option(command, "--generation-mode") == "conditional_prefix"
    assert option(command, "--conditioning-manifest") == str(MANIFEST)
    assert option(command, "--conditioning-manifest-sha256") == "a" * 64
    assert option(command, "--prefix-length") == "64"
    assert option(command, "--prompt-count") == "1024"
    assert option(command, "--diversity-prompt-count") == "256"
    assert option(command, "--completions-per-diversity-prompt") == "5"
    assert option(command, "--completion-schedule") == "c0:p0-1023;c1-4:p0-255"


def test_unconditional_command_does_not_emit_conditioning_flags() -> None:
    request = replace(conditional_request("mdlm"), generation_mode="unconditional")
    request = replace(
        request,
        conditioning_manifest=None,
        conditioning_manifest_sha256=None,
        conditioning_config_sha256=None,
        prefix_length=None,
        evaluation_continuation_length=None,
        prompt_count=None,
        diversity_prompt_count=None,
        completions_per_diversity_prompt=None,
        completion_schedule=None,
        results_root=None,
        sample_count=17,
    )

    command = MDLMAdapter().render_command(
        request,
        ROOT / "results" / "samples" / "lm1b" / "mdlm" / "steps_8",
        dry_run=True,
    )

    assert not any("conditioning" in argument or "generation-mode" in argument for argument in command)