"""Long-context baseline: the most recent visits that fit the token budget, oldest dropped first.

The budget is counted with the configured tokenizer (`answer.tokenizer` in backbone.yaml).
"""
from __future__ import annotations

import time

from .base import BaseAgent, MemoryBuildResult, MemoryContext


class LongContextAgent(BaseAgent):
    def __init__(self, *, max_context_tokens: int, margin: int, **params):
        """max_context_tokens is the input budget in tokens; margin is held back from it for the prompt and question."""
        super().__init__(**params)
        self.limit = max_context_tokens - margin
        self._visits: list[tuple[int, str, int]] = []

    def memorize(self, text: str, visit_number: int, date: str) -> MemoryBuildResult:
        t0 = time.time()
        self._visits.append((visit_number, text, self.count_tokens(text)))
        return MemoryBuildResult(success=True, method="long_context", action="append_visit",
                                 visit_number=visit_number, chunk_count=len(self._visits),
                                 added_chunks=1, time_cost=time.time() - t0)

    def build_memory_context(self, question: str, as_of: str = "",
                             display_question: str | None = None) -> MemoryContext:
        total = sum(n for _, _, n in self._visits)
        dropped: set[int] = set()
        i = 0
        while total > self.limit and i < len(self._visits):
            dropped.add(self._visits[i][0])
            total -= self._visits[i][2]
            i += 1
        mem = "\n\n".join(t for _, t, _ in self._visits[i:])
        return MemoryContext(
            mem=mem, n_tok=total,
            retrieved_visits=[v for v, _, _ in self._visits[i:]],
            context_truncated=bool(dropped),
            extra={"visits_in_mem": len(self._visits) - i, "visits_dropped": i},
        )