"""Time-based indexing for episodes."""

from __future__ import annotations

from bisect import bisect_left, bisect_right
from dataclasses import dataclass
from datetime import datetime
from typing import Iterable, List, Optional, Tuple


@dataclass(frozen=True)
class TemporalIndexEntry:
    timestamp: datetime
    episode_id: str
    counter: int


class TemporalIndex:
    """Maintain a sorted index of episode timestamps for fast temporal queries."""

    def __init__(self) -> None:
        self._entries: List[Tuple[datetime, int, str]] = []
        self._by_id: dict[str, Tuple[datetime, int]] = {}
        self._counter: int = 0

    def clear(self) -> None:
        self._entries.clear()
        self._by_id.clear()
        self._counter = 0

    def build(self, episodes: Iterable[TemporalIndexEntry]) -> None:
        self.clear()
        for entry in episodes:
            self.add(entry.episode_id, entry.timestamp, counter=entry.counter)

    def add(self, episode_id: str, timestamp: datetime, *, counter: Optional[int] = None) -> None:
        if episode_id in self._by_id:
            self.remove(episode_id)
        if counter is None:
            self._counter += 1
            counter = self._counter
        else:
            self._counter = max(self._counter, counter)
        key = (timestamp, counter, episode_id)
        idx = bisect_right(self._entries, key)
        self._entries.insert(idx, key)
        self._by_id[episode_id] = (timestamp, counter)

    def remove(self, episode_id: str) -> None:
        key = self._by_id.pop(episode_id, None)
        if key is None:
            return
        timestamp, counter = key
        needle = (timestamp, counter, episode_id)
        idx = bisect_left(self._entries, needle)
        if idx < len(self._entries) and self._entries[idx] == needle:
            self._entries.pop(idx)

    def query_range(
        self,
        start: Optional[datetime],
        end: Optional[datetime],
        *,
        limit: Optional[int] = None,
        ascending: bool = True,
    ) -> List[str]:
        if not self._entries:
            return []
        start_key = (start if start is not None else datetime.min, -1, "")
        end_key = (end if end is not None else datetime.max, self._counter + 1, "")
        left = bisect_left(self._entries, start_key)
        right = bisect_right(self._entries, end_key)
        entries = self._entries[left:right]
        if not ascending:
            entries = list(reversed(entries))
        ids = [entry[2] for entry in entries]
        if limit is not None and limit >= 0:
            return ids[:limit]
        return ids

    def before(
        self,
        timestamp: datetime,
        *,
        limit: Optional[int] = None,
        inclusive: bool = False,
    ) -> List[str]:
        cutoff = (timestamp, self._counter + 1 if inclusive else -1, "")
        right = bisect_right(self._entries, cutoff)
        ids = [entry[2] for entry in self._entries[:right]]
        if limit is not None and limit >= 0:
            return ids[-limit:]
        return ids

    def after(
        self,
        timestamp: datetime,
        *,
        limit: Optional[int] = None,
        inclusive: bool = False,
    ) -> List[str]:
        cutoff = (timestamp, -1 if inclusive else self._counter + 1, "")
        left = bisect_left(self._entries, cutoff)
        ids = [entry[2] for entry in self._entries[left:]]
        if limit is not None and limit >= 0:
            return ids[:limit]
        return ids
