"""Regression tests for FrozenModelBackend.format_prompt's chat-template handling."""

import unittest

from owmi.backends import FrozenModelBackend
from owmi.types import ExperimentConfig


def _make_config(system_prompt="You are a careful introspection assistant."):
    return ExperimentConfig(model_name="fake/model", system_prompt=system_prompt)


class _AcceptsSystemRoleTokenizer:
    """Mimics a normal chat template that accepts a dedicated system turn."""

    def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=True):
        assert not tokenize
        assert add_generation_prompt
        rendered = []
        for m in messages:
            rendered.append(f"<{m['role']}>{m['content']}")
        return "\n".join(rendered) + "\n<assistant>"


class _RejectsSystemRoleTokenizer:
    """Mimics Gemma-2's chat template, which raises when given a system turn."""

    def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=True):
        roles = [m["role"] for m in messages]
        if "system" in roles:
            raise Exception("System role not supported")
        rendered = []
        for m in messages:
            rendered.append(f"<{m['role']}>{m['content']}")
        return "\n".join(rendered) + "\n<assistant>"


class _RaisesUnrelatedErrorTokenizer:
    """A tokenizer whose apply_chat_template fails for an unrelated reason."""

    def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=True):
        raise Exception("tokenizer is not fully initialized")


def _backend_with_tokenizer(tokenizer, system_prompt="You are a careful introspection assistant."):
    backend = FrozenModelBackend.__new__(FrozenModelBackend)
    backend.tokenizer = tokenizer
    backend.config = _make_config(system_prompt)
    return backend


class FormatPromptSystemRoleTests(unittest.TestCase):
    def test_uses_system_role_when_supported(self):
        backend = _backend_with_tokenizer(_AcceptsSystemRoleTokenizer())
        out = backend.format_prompt("What happened?")
        self.assertIn("<system>You are a careful introspection assistant.", out)
        self.assertIn("<user>What happened?", out)

    def test_falls_back_to_merged_user_turn_when_system_role_rejected(self):
        backend = _backend_with_tokenizer(_RejectsSystemRoleTokenizer())
        out = backend.format_prompt("What happened?")
        # No system role in the final rendering; content survives, merged into the user turn.
        self.assertNotIn("<system>", out)
        self.assertIn("You are a careful introspection assistant.", out)
        self.assertIn("What happened?", out)
        self.assertIn("<user>", out)

    def test_unrelated_tokenizer_errors_still_raise(self):
        backend = _backend_with_tokenizer(_RaisesUnrelatedErrorTokenizer())
        with self.assertRaises(Exception) as ctx:
            backend.format_prompt("What happened?")
        self.assertIn("not fully initialized", str(ctx.exception))

    def test_no_tokenizer_falls_back_to_plain_format(self):
        backend = _backend_with_tokenizer(None)
        out = backend.format_prompt("What happened?")
        self.assertEqual(
            out,
            "System: You are a careful introspection assistant.\n"
            "User: What happened?\nAssistant:",
        )