from __future__ import annotations

from typing import Any, Callable, Optional

from contexthub.priority import select_by_priority
from contexthub.tokenizer import get_tokenizer
from contexthub.types import ContextItem, Priority


class ContextWindow:
    def __init__(
        self, max_tokens: int = 4096, tokenizer: Optional[Callable[[str], int]] = None
    ):
        if max_tokens <= 0:
            raise ValueError("max_tokens must be > 0")
        self.max_tokens = max_tokens
        self._tokenizer = tokenizer or get_tokenizer()
        self._items: list[ContextItem] = []

    def add(
        self,
        content: str,
        priority: Priority = Priority.MEDIUM,
        role: str = "user",
        metadata: Optional[dict[str, Any]] = None,
    ) -> ContextItem:
        item = ContextItem(
            content=content,
            priority=priority,
            role=role,
            metadata=metadata or {},
            token_count=self._tokenizer(content),
        )
        self._items.append(item)
        return item

    def build(self) -> list[ContextItem]:
        return select_by_priority(
            self._items, self.max_tokens, tokenizer=self._tokenizer
        )

    def to_messages(self) -> list[dict]:
        return [{"role": item.role, "content": item.content} for item in self.build()]

    def clear(self) -> None:
        self._items.clear()

    @property
    def total_tokens(self) -> int:
        total = 0
        for item in self._items:
            if item.token_count is None:
                item.token_count = self._tokenizer(item.content)
            total += item.token_count
        return total

    @property
    def remaining_tokens(self) -> int:
        return max(0, self.max_tokens - self.total_tokens)

    @property
    def items(self) -> list[ContextItem]:
        return list(self._items)
