"""Core context engine — the main entry point for the library."""

from __future__ import annotations

from typing import TYPE_CHECKING

from context_manager.budget import get_counter
from context_manager.models import ContextBlock
from context_manager.strategies import PriorityPruning

if TYPE_CHECKING:
    from context_manager.budget.base import TokenCounter
    from context_manager.distill.base import Compressor
    from context_manager.prune.deduplicator import Deduplicator
    from context_manager.strategies.base import PruningStrategy


class ContextEngine:
    """Build, budget, and compile an optimized context payload.

    Args:
        model: LLM model identifier (e.g. 'gpt-4', 'meta-llama/...').
        token_limit: Maximum tokens allowed in the compiled context.
        pruning_strategy: Strategy used to shed blocks when over budget.
        counter: Explicit TokenCounter; auto-resolved from *model* if None.
    """

    def __init__(
        self,
        model: str = "gpt-4",
        token_limit: int = 8192,
        pruning_strategy: PruningStrategy | None = None,
        counter: TokenCounter | None = None,
        deduplicator: Deduplicator | None = None,
        compressor: Compressor | None = None,
    ) -> None:
        self._model = model
        self._token_limit = token_limit
        self._strategy = pruning_strategy or PriorityPruning()
        self._counter = counter or get_counter(model)
        self._deduplicator = deduplicator
        self._compressor = compressor
        self._blocks: list[ContextBlock] = []
        self._compiled: list[ContextBlock] | None = None

    # -- Public API -----------------------------------------------------------

    def add(self, block: ContextBlock) -> ContextEngine:
        """Append a context block and count its tokens.

        Returns *self* for method chaining.
        """
        block.token_count = self._counter.count(block.content)
        self._blocks.append(block)
        self._compiled = None  # invalidate cache
        return self

    def compile(self) -> list[dict[str, str]]:
        """Prune and return the context as a list of message dicts.

        Each dict has keys ``role`` and ``content``, compatible with the
        OpenAI chat format. Framework-agnostic consumers can map this to
        their own message types trivially.

        When a :class:`Deduplicator` is configured, semantic
        deduplication runs **before** the pruning strategy.
        """
        candidates = list(self._blocks)

        # 1. Distill (Compress) flagged blocks
        if self._compressor is not None:
            # We must recreate the list because we might modify block content.
            # ContextBlock is a dataclass, so we use replace().
            import dataclasses

            new_candidates = []
            for block in candidates:
                if block.can_compress and block.content.strip():
                    compressed_content = self._compressor.compress(
                        block.content
                    )
                    # Create a copy with new content
                    new_block = dataclasses.replace(
                        block, content=compressed_content
                    )
                    new_candidates.append(new_block)
                else:
                    new_candidates.append(block)
            candidates = new_candidates

        # 2. Deduplicate
        if self._deduplicator is not None:
            candidates = self._deduplicator.deduplicate(candidates)

        # 3. Prune
        self._compiled = self._strategy.prune(candidates, self._token_limit)
        return [
            {"role": b.role, "content": b.content} for b in self._compiled
        ]

    def reset(self) -> None:
        """Clear all blocks and cached state."""
        self._blocks.clear()
        self._compiled = None

    # -- Introspection --------------------------------------------------------

    @property
    def blocks(self) -> list[ContextBlock]:
        """All blocks added so far (unfiltered)."""
        return list(self._blocks)

    @property
    def total_tokens(self) -> int:
        """Sum of token counts across all added blocks."""
        return sum(b.token_count or 0 for b in self._blocks)

    @property
    def compiled_tokens(self) -> int:
        """Token count of the last compiled output (0 if not compiled)."""
        if self._compiled is None:
            return 0
        return sum(b.token_count or 0 for b in self._compiled)

    @property
    def token_limit(self) -> int:
        """The configured token budget."""
        return self._token_limit

    @property
    def remaining_tokens(self) -> int:
        """Tokens still available in the budget after compilation."""
        return self._token_limit - self.compiled_tokens

    def summary(self) -> dict:
        """Return a diagnostic summary of the current state."""
        compiled = self._compiled or []
        dropped = [b for b in self._blocks if b not in compiled]
        return {
            "model": self._model,
            "token_limit": self._token_limit,
            "total_blocks": len(self._blocks),
            "kept_blocks": len(compiled),
            "dropped_blocks": len(dropped),
            "total_tokens": self.total_tokens,
            "compiled_tokens": self.compiled_tokens,
            "remaining_tokens": self.remaining_tokens,
            "dropped": [
                {"id": b.id, "role": b.role, "priority": b.priority.name, "tokens": b.token_count}
                for b in dropped
            ],
        }

    def visualize(self) -> str:
        """Print a simple text bar chart of token usage and return it."""
        compiled = self._compiled or []
        bar_width = 50
        limit = self._token_limit or 1

        lines: list[str] = [f"Token Budget: {self.compiled_tokens}/{limit}"]

        # Group by role
        role_tokens: dict[str, int] = {}
        for b in compiled:
            role_tokens[b.role] = role_tokens.get(b.role, 0) + (b.token_count or 0)

        for role, tokens in role_tokens.items():
            filled = int((tokens / limit) * bar_width)
            pct = (tokens / limit) * 100
            bar = "█" * filled + "░" * (bar_width - filled)
            lines.append(f"  {role:15s} [{bar}] {tokens:>5d} ({pct:.0f}%)")

        unused = limit - self.compiled_tokens
        filled = int((unused / limit) * bar_width)
        bar = "·" * filled + " " * (bar_width - filled)
        lines.append(f"  {'unused':15s} [{bar}] {unused:>5d} ({(unused / limit) * 100:.0f}%)")

        output = "\n".join(lines)
        print(output)
        return output
