from __future__ import annotations

from typing import Callable, Optional

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


def trim_to_fit(
    text: str,
    max_tokens: int,
    tokenizer: Optional[Callable[[str], int]] = None,
    suffix: str = "...",
) -> str:
    if max_tokens <= 0:
        return ""
    if not text:
        return ""

    count = tokenizer or get_tokenizer()
    if count(text) <= max_tokens:
        return text

    if count(suffix) >= max_tokens:
        return ""

    words = text.split()
    result_words: list[str] = []
    for word in words:
        candidate = " ".join(result_words + [word]).strip()
        candidate_with_suffix = f"{candidate}{suffix}"
        if count(candidate_with_suffix) <= max_tokens:
            result_words.append(word)
        else:
            break

    if not result_words:
        return ""
    return f"{' '.join(result_words)}{suffix}"


def select_by_priority(
    items: list[ContextItem],
    max_tokens: int,
    tokenizer: Optional[Callable[[str], int]] = None,
) -> list[ContextItem]:
    if max_tokens <= 0:
        return []

    count = tokenizer or get_tokenizer()

    annotated = list(enumerate(items))
    ordered = sorted(annotated, key=lambda pair: (pair[1].priority.value, pair[0]))

    selected: list[ContextItem] = []
    total = 0

    for _, item in ordered:
        tokens = (
            item.token_count if item.token_count is not None else count(item.content)
        )
        item.token_count = tokens
        if item.priority == Priority.CRITICAL:
            selected.append(item)
            total += tokens
            continue

        if total + tokens <= max_tokens:
            selected.append(item)
            total += tokens

    return selected
