from __future__ import annotations

import time
from datetime import datetime
from typing import Iterable

import yfinance as yf

from .models import Tick
from .provider import TickProvider


class YahooPollingProvider(TickProvider):
    def __init__(
        self,
        symbols: list[str],
        market: str = "BEI",
        poll_seconds: int = 60,
        bootstrap_bars: int = 60
    ) -> None:
        self.symbols = symbols
        self.market = market
        self.poll_seconds = poll_seconds
        self.bootstrap_bars = bootstrap_bars
        self._last_timestamp: dict[str, datetime] = {}
        self._bootstrapped = False

    def _download(self):
        return yf.download(
            self.symbols,
            period="5d",
            interval="1m",
            progress=False,
            group_by="column",
            threads=True,
        )

    def _ticks_for_symbol(self, data, symbol: str, bootstrap: bool):
        close = data["Close"][symbol].dropna()
        volume = data["Volume"][symbol].fillna(0)

        if close.empty:
            return

        indexes = close.index[-self.bootstrap_bars:] if bootstrap else close.index[-1:]

        for index in indexes:
            timestamp = index.to_pydatetime()

            if self._last_timestamp.get(symbol) == timestamp:
                continue

            price = float(close.loc[index])
            cumulative_volume = float(volume.loc[:index].sum())

            self._last_timestamp[symbol] = timestamp

            yield Tick(
                symbol=symbol.replace(".JK", ""),
                market=self.market,
                price=price,
                volume=cumulative_volume,
                timestamp=timestamp,
                source="YAHOO",
            )

    def stream(self) -> Iterable[Tick]:
        while True:
            try:
                data = self._download()
                bootstrap = not self._bootstrapped

                for symbol in self.symbols:
                    try:
                        yield from self._ticks_for_symbol(
                            data,
                            symbol,
                            bootstrap=bootstrap,
                        )
                    except Exception as exc:
                        print({
                            "event": "yahoo-symbol-error",
                            "symbol": symbol,
                            "error": str(exc),
                        })

                self._bootstrapped = True

            except Exception as exc:
                print({
                    "event": "yahoo-provider-error",
                    "error": str(exc),
                })

            time.sleep(self.poll_seconds)