"""Tests for prompt_builder.compiler

These are the most important tests. The compiler is the core of the framework.
"""

import pytest
from prompt_builder import (
    Compiler,
    Prompt,
    Instruction,
    Persona,
    Rule,
    ContextBlock,
    Example,
    InputData,
    OutputSpec,
    CoTSpec,
    CoTStrategy,
    Visibility,
    Message,
    TokenBudget,
    PriorityMap,
    PriorityEntry,
    Section,
    TruncationStrategy,
    compile,
)
from prompt_builder.tokens import TokenCounter


@pytest.fixture
def counter():
    return TokenCounter()


@pytest.fixture
def compiler(counter):
    return Compiler(counter)


# ── Basic compilation ──────────────────────────────────────────────────────


class TestBasicCompilation:
    def test_minimal_prompt(self, compiler):
        """Simplest possible prompt compiles to a user message."""
        prompt = Prompt(instruction=Instruction(task="Say hello."))
        budget = TokenBudget(context_window=8000)

        result = compiler.compile(prompt, budget)

        assert len(result.messages) >= 1
        # Should have a user message with the instruction
        user_msgs = [m for m in result.messages if m["role"] == "user"]
        assert len(user_msgs) == 1
        assert "Say hello." in user_msgs[0]["content"]

    def test_with_persona_creates_system_message(self, compiler):
        """Persona goes into the system message."""
        prompt = Prompt(
            persona=Persona(role="You are a pirate."),
            instruction=Instruction(task="Greet me."),
        )
        budget = TokenBudget(context_window=8000)

        result = compiler.compile(prompt, budget)

        sys_msgs = [m for m in result.messages if m["role"] == "system"]
        assert len(sys_msgs) == 1
        assert "pirate" in sys_msgs[0]["content"]

    def test_message_ordering(self, compiler):
        """Messages come in order: system, history, user."""
        prompt = Prompt(
            persona=Persona(role="You are helpful."),
            instruction=Instruction(task="Do it."),
            history=[
                Message(role="user", content="Hi"),
                Message(role="assistant", content="Hello"),
            ],
        )
        budget = TokenBudget(context_window=8000)

        result = compiler.compile(prompt, budget)

        roles = [m["role"] for m in result.messages]
        assert roles[0] == "system"
        # History comes between system and final user message
        assert "user" in roles[1:-1] or "assistant" in roles[1:-1]
        assert roles[-1] == "user"

    def test_all_sections_render(self, compiler):
        """Prompt with all sections populated compiles without error."""
        prompt = Prompt(
            persona=Persona(role="Expert.", tone="formal", domain=["ml"]),
            rules=[Rule(constraint="Be precise.")],
            context=[ContextBlock(content="Relevant fact.", source="rag", relevance=0.9)],
            examples=[Example(input="q", output="a")],
            instruction=Instruction(task="Classify.", intent="Routing."),
            input_data=InputData(content="ticket text"),
            output_spec=OutputSpec(format="json"),
            cot=CoTSpec(strategy=CoTStrategy.STEP_BY_STEP, visibility=Visibility.VISIBLE),
            history=[Message(role="user", content="Previous turn")],
        )
        budget = TokenBudget(context_window=128_000)

        result = compiler.compile(prompt, budget)

        assert len(result.messages) >= 3  # system + history + user
        assert result.total_tokens > 0
        assert result.headroom > 0
        assert len(result.warnings) == 0  # should all fit in 128k


class TestTokenCounting:
    def test_total_tokens_matches_messages(self, compiler, counter):
        """Reported tokens should match counting the actual messages."""
        prompt = Prompt(
            persona=Persona(role="You are a helpful assistant."),
            instruction=Instruction(task="Help me."),
        )
        budget = TokenBudget(context_window=8000)

        result = compiler.compile(prompt, budget)

        # Total should be close to counting the messages directly
        direct_count = counter.count_messages(result.messages)
        assert result.total_tokens == direct_count

    def test_headroom_calculation(self, compiler):
        prompt = Prompt(instruction=Instruction(task="X"))
        budget = TokenBudget(context_window=8000, output_reserve=2000, safety_margin=100)

        result = compiler.compile(prompt, budget)

        assert result.budget_available == 8000 - 2000 - 100
        assert result.headroom == result.budget_available - result.total_tokens


# ── Truncation ─────────────────────────────────────────────────────────────


class TestTruncation:
    def _make_fat_prompt(self, n_context=10, content_size=500):
        """Helper: create a prompt with lots of context blocks."""
        context = []
        for i in range(n_context):
            context.append(ContextBlock(
                content=f"Context block {i}. " + ("x " * content_size),
                source="test",
                relevance=i / n_context,  # 0.0 to 0.9
            ))
        return Prompt(
            instruction=Instruction(task="Summarize all context."),
            context=context,
        )

    def test_no_truncation_when_fits(self, compiler):
        """No warnings when everything fits."""
        prompt = Prompt(instruction=Instruction(task="Hello."))
        budget = TokenBudget(context_window=8000)

        result = compiler.compile(prompt, budget)
        assert len(result.warnings) == 0

    def test_context_dropped_by_relevance(self, compiler):
        """Lowest-relevance context blocks get dropped first."""
        prompt = self._make_fat_prompt(n_context=10, content_size=200)
        # Tight budget — can't fit all 10 blocks
        budget = TokenBudget(context_window=3000, output_reserve=700)

        result = compiler.compile(prompt, budget)

        # Should have warnings about dropping
        assert len(result.warnings) > 0
        assert any("Dropped context" in w for w in result.warnings)

        # The context section report should show items dropped
        ctx_report = next(s for s in result.sections if s.section == Section.CONTEXT)
        assert ctx_report.items_dropped > 0
        assert ctx_report.items_included > 0
        assert ctx_report.items_included < 10

    def test_lowest_relevance_dropped_first(self, compiler):
        """Verify drop order: lowest relevance goes first."""
        # 3 blocks: relevance 0.1, 0.5, 0.9
        prompt = Prompt(
            instruction=Instruction(task="Summarize."),
            context=[
                ContextBlock(content="Low relevance. " + ("x " * 500), relevance=0.1, source="a"),
                ContextBlock(content="Mid relevance. " + ("x " * 500), relevance=0.5, source="b"),
                ContextBlock(content="High relevance. " + ("x " * 500), relevance=0.9, source="c"),
            ],
        )
        # Budget tight enough to drop 1-2 blocks
        budget = TokenBudget(context_window=1900, output_reserve=200)

        result = compiler.compile(prompt, budget)

        # Check user message: high relevance should be present
        user_content = next(m["content"] for m in result.messages if m["role"] == "user")
        assert "High relevance." in user_content

        # Low relevance should have been dropped (in warnings)
        dropped_warnings = [w for w in result.warnings if "sort_key=0.10" in w]
        assert len(dropped_warnings) > 0

    def test_history_truncated_oldest_first(self, compiler):
        """Oldest history messages get dropped first."""
        history = [
            Message(role="user", content=f"Turn {i}. " + ("x " * 200))
            for i in range(10)
        ]
        prompt = Prompt(
            instruction=Instruction(task="Continue."),
            history=history,
        )
        budget = TokenBudget(context_window=2300, output_reserve=100)

        result = compiler.compile(prompt, budget)

        # Most recent turns should be present
        all_content = " ".join(m["content"] for m in result.messages)
        assert "Turn 9" in all_content  # most recent should survive

        # Oldest should be dropped
        hist_report = next(s for s in result.sections if s.section == Section.HISTORY)
        assert hist_report.items_dropped > 0

    def test_examples_dropped_by_weight(self, compiler):
        """Lowest-weight examples get dropped first."""
        examples = [
            Example(input="q0", output="a0 " + ("x " * 200), weight=0.1),
            Example(input="q1", output="a1 " + ("x " * 200), weight=0.5),
            Example(input="q2", output="a2 " + ("x " * 200), weight=0.9),
        ]
        prompt = Prompt(
            instruction=Instruction(task="Answer."),
            examples=examples,
        )
        budget = TokenBudget(context_window=2000, output_reserve=500)

        result = compiler.compile(prompt, budget)

        user_content = next(m["content"] for m in result.messages if m["role"] == "user")
        # Highest weight example should survive
        assert "a2" in user_content

    def test_never_cut_sections_survive(self, compiler):
        """Instruction and persona are never cut, even under pressure."""
        prompt = Prompt(
            persona=Persona(role="You are critical."),
            instruction=Instruction(task="This must survive."),
            context=[
                ContextBlock(content="Expendable. " + ("x " * 2000), relevance=0.1, source="test"),
            ],
        )
        budget = TokenBudget(context_window=500, output_reserve=100)

        result = compiler.compile(prompt, budget)

        sys_content = next(
            (m["content"] for m in result.messages if m["role"] == "system"), ""
        )
        user_content = next(
            (m["content"] for m in result.messages if m["role"] == "user"), ""
        )
        assert "critical" in sys_content
        assert "must survive" in user_content

    def test_cot_cut_last(self, compiler):
        """CoT should be cut only after context/history/examples."""
        prompt = Prompt(
            instruction=Instruction(task="Think."),
            cot=CoTSpec(strategy=CoTStrategy.STEP_BY_STEP, visibility=Visibility.VISIBLE),
            context=[
                ContextBlock(content="Data. " + ("x " * 500), relevance=0.3, source="test"),
            ],
        )
        # Budget that requires cutting something
        budget = TokenBudget(context_window=800, output_reserve=200)

        result = compiler.compile(prompt, budget)

        # Context should be cut before CoT
        drop_warnings = [w for w in result.warnings if "Dropped" in w or "Cut" in w]
        if drop_warnings:
            context_cut = any("context" in w for w in drop_warnings)
            cot_cut = any("cot" in w for w in drop_warnings)
            if context_cut and not cot_cut:
                pass  # correct: context cut, cot preserved
            # If both are cut, that's fine too (budget really tight)


class TestCustomPriority:
    def test_custom_priority_map(self, compiler):
        """Custom priority map changes what gets cut."""
        prompt = Prompt(
            instruction=Instruction(task="Go."),
            context=[
                ContextBlock(content="Keep this. " + ("x " * 300), relevance=0.1, source="a"),
            ],
            examples=[
                Example(input="q", output="a " + ("x " * 300), weight=0.9),
            ],
            # Custom: cut examples before context (reverse of default)
            priority=PriorityMap(entries=[
                PriorityEntry(section=Section.EXAMPLES, strategy=TruncationStrategy.CUT_BY_WEIGHT),
                PriorityEntry(section=Section.CONTEXT, strategy=TruncationStrategy.CUT_BY_RELEVANCE),
                PriorityEntry(section=Section.INSTRUCTION, strategy=TruncationStrategy.NEVER_CUT),
            ]),
        )
        budget = TokenBudget(context_window=1000, output_reserve=200)

        result = compiler.compile(prompt, budget)

        # Examples should be cut before context
        if result.warnings:
            first_cut = result.warnings[1] if len(result.warnings) > 1 else result.warnings[0]
            # At minimum, the compilation should succeed
            assert result.messages is not None


# ── Compilation Report ─────────────────────────────────────────────────────


class TestCompilationReport:
    def test_report_has_all_present_sections(self, compiler):
        """Report includes entries for every section in the prompt."""
        prompt = Prompt(
            persona=Persona(role="Helper."),
            instruction=Instruction(task="Do it."),
            context=[ContextBlock(content="Info.", source="test")],
            examples=[Example(input="q", output="a")],
        )
        budget = TokenBudget(context_window=128_000)

        result = compiler.compile(prompt, budget)

        section_names = {s.section for s in result.sections}
        assert Section.PERSONA in section_names
        assert Section.INSTRUCTION in section_names
        assert Section.CONTEXT in section_names
        assert Section.EXAMPLES in section_names

    def test_report_string(self, compiler):
        prompt = Prompt(instruction=Instruction(task="Test."))
        budget = TokenBudget(context_window=8000)

        result = compiler.compile(prompt, budget)

        report = result.report()
        assert isinstance(report, str)
        assert "tokens" in report


# ── Convenience function ───────────────────────────────────────────────────


class TestConvenienceFunction:
    def test_compile_function(self):
        """The module-level compile() function works."""
        prompt = Prompt(instruction=Instruction(task="Hello."))
        budget = TokenBudget(context_window=8000)

        result = compile(prompt, budget)

        assert len(result.messages) >= 1
        assert result.total_tokens > 0
