from __future__ import annotations

import tempfile
import unittest
from pathlib import Path

from hcl.evaluator import EvaluationRecorder
from hcl.memory import MemoryConfig
from hcl.optimizer import Optimizer
from hcl.pipeline import HCLPipeline


class _Model:
    def generate(self, prompt: str, *, state: dict[str, object]) -> str:
        del prompt, state
        return "{}"


class DistillationModelRoutingTest(unittest.TestCase):
    def _pipeline(
        self,
        directory: str,
        *,
        model: _Model,
        judge_model: _Model | None = None,
        selection_model: _Model | None = None,
        memory_model: _Model | None = None,
        judge_memory: bool = False,
    ) -> HCLPipeline:
        return HCLPipeline(
            model,
            judge_model=judge_model,
            selection_model=selection_model,
            memory_model=memory_model,
            judge_memory=judge_memory,
            evaluation_recorder=EvaluationRecorder(Path(directory) / "evaluator"),
            optimizer=Optimizer(record_dir=Path(directory) / "optimizer", candidate_count=0),
            memory_config=MemoryConfig(enabled=False, record_dir=Path(directory) / "memory"),
        )

    def test_judge_defaults_to_main_model(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            model = _Model()
            pipeline = self._pipeline(directory, model=model)

        self.assertIs(pipeline.judge_model, model)
        self.assertIs(pipeline.selection_model, model)
        self.assertIs(pipeline.memory_model, model)
        self.assertIs(pipeline.memory.model, model)

    def test_partial_distillation_keeps_memory_on_main_model(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            model = _Model()
            judge = _Model()
            pipeline = self._pipeline(directory, model=model, judge_model=judge)

        self.assertIs(pipeline.judge_model, judge)
        self.assertIs(pipeline.memory.model, model)

    def test_full_distillation_routes_memory_to_judge(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            model = _Model()
            judge = _Model()
            pipeline = self._pipeline(
                directory,
                model=model,
                judge_model=judge,
                judge_memory=True,
            )

        self.assertIs(pipeline.judge_model, judge)
        self.assertIs(pipeline.memory.model, judge)

    def test_online_teacher_routes_selection_and_memory_but_not_final_or_optimizer(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            model = _Model()
            teacher = _Model()
            pipeline = self._pipeline(
                directory,
                model=model,
                selection_model=teacher,
                memory_model=teacher,
            )

        self.assertIs(pipeline.model, model)
        self.assertIs(pipeline.judge_model, model)
        self.assertIs(pipeline.selection_model, teacher)
        self.assertIs(pipeline.memory_model, teacher)
        self.assertIs(pipeline.memory.model, teacher)


if __name__ == "__main__":
    unittest.main()