from collections import defaultdict, deque
from typing import Protocol

from app.bars.models import Bar, Timeframe


class BarStore(Protocol):
    def put(self, bar: Bar) -> None: ...

    def get(self, tf: Timeframe, limit: int | None = None) -> list[Bar]: ...


class InMemoryBarStore:
    def __init__(self, max_bars_per_tf: int = 5000):
        self._bars: dict[Timeframe, deque[Bar]] = defaultdict(
            lambda: deque(maxlen=max_bars_per_tf)
        )

    def put(self, bar: Bar) -> None:
        bars = self._bars[bar.tf]
        if bars and bars[-1].t == bar.t:
            bars[-1] = bar
        elif not bars or bar.t > bars[-1].t:
            bars.append(bar)

    def get(self, tf: Timeframe, limit: int | None = None) -> list[Bar]:
        bars = list(self._bars[tf])
        return bars[-limit:] if limit is not None else bars

    def counts(self) -> dict[str, int]:
        return {tf.value: len(self._bars[tf]) for tf in Timeframe}