"""Compiler: token-aware prompt assembly with priority-based truncation.

This is the core of the framework. Takes a Prompt + TokenBudget and produces
a CompilationResult with the final message array and a full report of what
was included, dropped, and why.

Phases:
  1. Measure — compute token costs for every renderable section
  2. Check — does everything fit?
  3. Trim — apply priority map to shed tokens if over budget
  4. Render — build the final message array
"""

from __future__ import annotations

from dataclasses import dataclass, field

from . import renderer
from .tokens import TokenCounter, get_counter
from .types import (
    CompilationResult,
    ContextBlock,
    CoTSpec,
    CoTStrategy,
    Example,
    Message,
    Prompt,
    Section,
    SectionReport,
    TokenBudget,
    TruncationStrategy,
)


@dataclass
class _MeasuredSection:
    """Internal: a section with its rendered text and token cost."""

    section: Section
    text: str
    tokens: int
    # For list-type sections, individual items with costs
    items: list[_MeasuredItem] = field(default_factory=list)


@dataclass
class _MeasuredItem:
    """Internal: a single item within a list-type section."""

    index: int
    text: str
    tokens: int
    sort_key: float  # relevance for context, weight for examples, timestamp for history
    included: bool = True


class Compiler:
    """Compiles a Prompt into a message array within a token budget.

    Usage:
        compiler = Compiler()
        result = compiler.compile(prompt, budget)
        print(result.report())
    """

    def __init__(self, counter: TokenCounter | None = None) -> None:
        self.counter = counter or get_counter()

    def compile(self, prompt: Prompt, budget: TokenBudget) -> CompilationResult:
        """Main entry point. Compile prompt into messages within budget."""

        # Phase 1: Measure
        measured = self._measure(prompt)

        # Phase 2: Check total
        total = sum(s.tokens for s in measured.values())
        available = budget.available

        # Phase 3: Trim if needed
        reports: list[SectionReport] = []
        warnings: list[str] = []

        if total > available:
            overage = total - available
            warnings.append(f"Over budget by {overage} tokens. Trimming.")
            self._trim(measured, prompt, available, warnings)

        # Phase 4: Render into messages
        messages, reports = self._render_messages(measured, prompt)

        final_tokens = self.counter.count_messages(messages)

        return CompilationResult(
            messages=messages,
            sections=reports,
            total_tokens=final_tokens,
            budget_available=available,
            warnings=warnings,
        )

    # ── Phase 1: Measure ───────────────────────────────────────────────────

    def _measure(self, prompt: Prompt) -> dict[Section, _MeasuredSection]:
        """Compute token costs for all sections."""
        measured: dict[Section, _MeasuredSection] = {}

        # Persona
        if prompt.persona:
            text = renderer.render_persona(prompt.persona)
            measured[Section.PERSONA] = _MeasuredSection(
                section=Section.PERSONA,
                text=text,
                tokens=self.counter.count(text),
            )

        # Rules
        if prompt.rules:
            text = renderer.render_rules(prompt.rules)
            measured[Section.RULES] = _MeasuredSection(
                section=Section.RULES,
                text=text,
                tokens=self.counter.count(text),
            )

        # Output spec
        if prompt.output_spec:
            text = renderer.render_output_spec(prompt.output_spec)
            if text:
                measured[Section.OUTPUT_SPEC] = _MeasuredSection(
                    section=Section.OUTPUT_SPEC,
                    text=text,
                    tokens=self.counter.count(text),
                )

        # CoT
        if prompt.cot and prompt.cot.strategy != CoTStrategy.NONE:
            text = renderer.render_cot(prompt.cot)
            if text:
                measured[Section.COT] = _MeasuredSection(
                    section=Section.COT,
                    text=text,
                    tokens=self.counter.count(text),
                )

        # Context blocks (list — individually measurable)
        if prompt.context:
            items = []
            for i, block in enumerate(prompt.context):
                text = f'<document source="{block.source}">\n{block.content}\n</document>'
                tokens = block.token_count or self.counter.count(text)
                items.append(_MeasuredItem(
                    index=i,
                    text=text,
                    tokens=tokens,
                    sort_key=block.relevance,
                ))
            full_text = renderer.render_context_blocks(prompt.context)
            measured[Section.CONTEXT] = _MeasuredSection(
                section=Section.CONTEXT,
                text=full_text,
                tokens=sum(it.tokens for it in items) + 10,  # wrapper tags overhead
                items=items,
            )

        # Examples (list — individually measurable)
        if prompt.examples:
            items = []
            for i, ex in enumerate(prompt.examples):
                text = renderer.render_examples([ex])
                tokens = self.counter.count(text)
                items.append(_MeasuredItem(
                    index=i,
                    text=text,
                    tokens=tokens,
                    sort_key=ex.weight,
                ))
            full_text = renderer.render_examples(prompt.examples)
            measured[Section.EXAMPLES] = _MeasuredSection(
                section=Section.EXAMPLES,
                text=full_text,
                tokens=sum(it.tokens for it in items) + 5,  # header overhead
                items=items,
            )

        # History (list — truncated oldest first)
        if prompt.history:
            items = []
            for i, msg in enumerate(prompt.history):
                tokens = msg.token_count or self.counter.count(msg.content)
                items.append(_MeasuredItem(
                    index=i,
                    text=msg.content,
                    tokens=tokens + 4,  # role overhead
                    sort_key=float(i),  # lower index = older = cut first
                ))
            measured[Section.HISTORY] = _MeasuredSection(
                section=Section.HISTORY,
                text="",  # rendered differently (as separate messages)
                tokens=sum(it.tokens for it in items),
                items=items,
            )

        # Instruction
        text = renderer.render_instruction(prompt.instruction)
        measured[Section.INSTRUCTION] = _MeasuredSection(
            section=Section.INSTRUCTION,
            text=text,
            tokens=self.counter.count(text),
        )

        # Input data
        if prompt.input_data:
            text = renderer.render_input_data(prompt.input_data)
            measured[Section.INPUT_DATA] = _MeasuredSection(
                section=Section.INPUT_DATA,
                text=text,
                tokens=self.counter.count(text),
            )

        return measured

    # ── Phase 3: Trim ──────────────────────────────────────────────────────

    def _trim(
        self,
        measured: dict[Section, _MeasuredSection],
        prompt: Prompt,
        target: int,
        warnings: list[str],
    ) -> None:
        """Apply priority map to shed tokens until under target."""
        current = sum(s.tokens for s in measured.values())

        for entry in prompt.priority.cutting_order():
            if current <= target:
                break

            section = entry.section
            if section not in measured:
                continue

            ms = measured[section]
            strategy = entry.strategy

            if strategy == TruncationStrategy.NEVER_CUT:
                continue

            if strategy == TruncationStrategy.CUT_LAST:
                # Only cut if we've exhausted other options
                # (CUT_LAST sections are at the end of cutting_order)
                if current <= target:
                    break
                # Remove entirely
                tokens_freed = ms.tokens
                ms.tokens = 0
                ms.text = ""
                ms.items = []
                current -= tokens_freed
                warnings.append(
                    f"Cut entire {section.value} section ({tokens_freed} tokens)"
                )
                continue

            # For list-type sections, drop items one at a time
            if not ms.items:
                # Non-list section — drop entirely
                tokens_freed = ms.tokens
                ms.tokens = 0
                ms.text = ""
                current -= tokens_freed
                warnings.append(
                    f"Cut entire {section.value} section ({tokens_freed} tokens)"
                )
                continue

            # Sort items by sort_key to determine drop order
            if strategy == TruncationStrategy.CUT_BY_RELEVANCE:
                # Lowest relevance first
                drop_order = sorted(
                    [it for it in ms.items if it.included],
                    key=lambda it: it.sort_key,
                )
            elif strategy == TruncationStrategy.CUT_BY_WEIGHT:
                # Lowest weight first
                drop_order = sorted(
                    [it for it in ms.items if it.included],
                    key=lambda it: it.sort_key,
                )
            elif strategy == TruncationStrategy.TRUNCATE_OLDEST:
                # Oldest (lowest index) first
                drop_order = sorted(
                    [it for it in ms.items if it.included],
                    key=lambda it: it.sort_key,
                )
            else:
                continue

            for item in drop_order:
                if current <= target:
                    break
                item.included = False
                current -= item.tokens
                ms.tokens -= item.tokens
                warnings.append(
                    f"Dropped {section.value} item {item.index} "
                    f"({item.tokens} tokens, sort_key={item.sort_key:.2f})"
                )

    # ── Phase 4: Render ────────────────────────────────────────────────────

    def _render_messages(
        self,
        measured: dict[Section, _MeasuredSection],
        prompt: Prompt,
    ) -> tuple[list[dict[str, str]], list[SectionReport]]:
        """Build final message array and section reports."""

        messages: list[dict[str, str]] = []
        reports: list[SectionReport] = []

        # ── System message ──
        system_parts: list[str] = []

        for section in [Section.PERSONA, Section.RULES, Section.OUTPUT_SPEC, Section.COT]:
            if section in measured and measured[section].text:
                system_parts.append(measured[section].text)

            # Build report for this section
            if section in measured:
                ms = measured[section]
                reports.append(SectionReport(
                    section=section,
                    included_tokens=ms.tokens,
                    original_tokens=ms.tokens,  # non-list sections: included = original after trim
                    items_included=1 if ms.text else 0,
                    items_dropped=0,
                ))

        if system_parts:
            messages.append({
                "role": "system",
                "content": "\n\n".join(system_parts),
            })

        # ── History messages ──
        if Section.HISTORY in measured:
            ms = measured[Section.HISTORY]
            included_items = [it for it in ms.items if it.included]
            dropped_items = [it for it in ms.items if not it.included]

            for item in sorted(included_items, key=lambda it: it.index):
                history_msg = prompt.history[item.index]
                messages.append({
                    "role": history_msg.role,
                    "content": history_msg.content,
                })

            reports.append(SectionReport(
                section=Section.HISTORY,
                included_tokens=sum(it.tokens for it in included_items),
                original_tokens=sum(it.tokens for it in ms.items),
                items_included=len(included_items),
                items_dropped=len(dropped_items),
                strategy_applied=(
                    TruncationStrategy.TRUNCATE_OLDEST if dropped_items else None
                ),
            ))

        # ── User message ──
        user_parts: list[str] = []

        # Context blocks
        if Section.CONTEXT in measured:
            ms = measured[Section.CONTEXT]
            included_items = [it for it in ms.items if it.included]
            dropped_items = [it for it in ms.items if not it.included]

            if included_items:
                included_blocks = [
                    prompt.context[it.index]
                    for it in sorted(included_items, key=lambda x: -x.sort_key)
                ]
                user_parts.append(renderer.render_context_blocks(included_blocks))

            reports.append(SectionReport(
                section=Section.CONTEXT,
                included_tokens=sum(it.tokens for it in included_items),
                original_tokens=sum(it.tokens for it in ms.items),
                items_included=len(included_items),
                items_dropped=len(dropped_items),
                strategy_applied=(
                    TruncationStrategy.CUT_BY_RELEVANCE if dropped_items else None
                ),
            ))

        # Examples
        if Section.EXAMPLES in measured:
            ms = measured[Section.EXAMPLES]
            included_items = [it for it in ms.items if it.included]
            dropped_items = [it for it in ms.items if not it.included]

            if included_items:
                included_examples = [
                    prompt.examples[it.index]
                    for it in sorted(included_items, key=lambda x: -x.sort_key)
                ]
                user_parts.append(renderer.render_examples(included_examples))

            reports.append(SectionReport(
                section=Section.EXAMPLES,
                included_tokens=sum(it.tokens for it in included_items),
                original_tokens=sum(it.tokens for it in ms.items),
                items_included=len(included_items),
                items_dropped=len(dropped_items),
                strategy_applied=(
                    TruncationStrategy.CUT_BY_WEIGHT if dropped_items else None
                ),
            ))

        # Instruction (always present)
        if Section.INSTRUCTION in measured:
            ms = measured[Section.INSTRUCTION]
            user_parts.append(ms.text)
            reports.append(SectionReport(
                section=Section.INSTRUCTION,
                included_tokens=ms.tokens,
                original_tokens=ms.tokens,
                items_included=1,
                items_dropped=0,
            ))

        # Input data
        if Section.INPUT_DATA in measured:
            ms = measured[Section.INPUT_DATA]
            if ms.text:
                user_parts.append(ms.text)
            reports.append(SectionReport(
                section=Section.INPUT_DATA,
                included_tokens=ms.tokens,
                original_tokens=ms.tokens,
                items_included=1 if ms.text else 0,
                items_dropped=0,
            ))

        if user_parts:
            messages.append({
                "role": "user",
                "content": "\n\n".join(user_parts),
            })

        return messages, reports


def compile(prompt: Prompt, budget: TokenBudget, counter: TokenCounter | None = None) -> CompilationResult:
    """Convenience function. Compile a prompt within a token budget."""
    return Compiler(counter).compile(prompt, budget)
