from dataclasses import replace

from app.bars.models import Bar, Timeframe
from app.bars.session import bucket_start


class Aggregator:
    def __init__(self, timeframes: list[Timeframe] | None = None):
        self.timeframes = timeframes or list(Timeframe)
        self.forming: dict[Timeframe, Bar] = {}

    @staticmethod
    def _can_aggregate(source: Timeframe, target: Timeframe) -> bool:
        if source is Timeframe.M1:
            return True
        if source is Timeframe.H1:
            return target in (Timeframe.H1, Timeframe.D1)
        return source is target

    def update(self, incoming: Bar) -> list[Bar]:
        emitted: list[Bar] = []
        for tf in self.timeframes:
            if not self._can_aggregate(incoming.tf, tf):
                continue
            if tf is incoming.tf:
                emitted.append(replace(incoming))
                continue

            start = bucket_start(incoming.t, tf)
            current = self.forming.get(tf)
            if current is not None and start < current.t:
                continue
            if current is None or start > current.t:
                if current is not None:
                    emitted.append(replace(current, closed=True))
                current = Bar(
                    tf=tf,
                    t=start,
                    o=incoming.o,
                    h=incoming.h,
                    l=incoming.l,
                    c=incoming.c,
                    v=incoming.v,
                    closed=False,
                    symbol=incoming.symbol,
                    source=incoming.source,
                )
                self.forming[tf] = current
            else:
                current.h = max(current.h, incoming.h)
                current.l = min(current.l, incoming.l)
                current.c = incoming.c
                current.v += incoming.v
                current.symbol = incoming.symbol
                current.source = incoming.source
            emitted.append(replace(current))
        return emitted