from __future__ import annotations

import asyncio
import time
from dataclasses import dataclass
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, Dict, List, Tuple, Set

import pyarrow.parquet as pq

from app.config import settings
from app.streaming.live_store import live_store, _TIMEFRAME_SECONDS


@dataclass(frozen=True)
class ReplayKey:
    symbol: str
    timeframe: str


class HistoricalReplayProvider:
    """Streams OHLCV data by replaying parquet partitions as live candles."""

    def __init__(
        self,
        root: str | Path | None = None,
        speed: float | None = None,
        days: int | None = None,
        loop: bool | None = None,
    ) -> None:
        default_root = Path(__file__).resolve().parents[3] / "data" / "parquet" / "live"
        resolved_root = Path(root) if root else Path(settings.HISTORICAL_REPLAY_ROOT or default_root)
        self._root = resolved_root
        self._speed = max(speed or settings.HISTORICAL_REPLAY_SPEED, 0.1)
        self._days_back = max(days or settings.HISTORICAL_REPLAY_DAYS, 1)
        self._loop = settings.HISTORICAL_REPLAY_LOOP if loop is None else loop

        self._tasks: Dict[ReplayKey, asyncio.Task] = {}
        self._pinned: Set[ReplayKey] = set()
        self._lock = asyncio.Lock()
        self._active_counts: Dict[ReplayKey, int] = {}
        self._bar_cache: Dict[ReplayKey, Tuple[float, List[Dict[str, Any]]]] = {}
        self._positions: Dict[ReplayKey, int] = {}
        self._last_times: Dict[ReplayKey, int] = {}

    def _sleep_seconds(self, timeframe: str) -> float:
        base = _TIMEFRAME_SECONDS.get(timeframe.lower(), 60)
        return max(base / self._speed, 0.5)

    def _ensure_history(self, key: ReplayKey) -> None:
        try:
            if not live_store.get_history(key.symbol, key.timeframe):
                live_store.load_historical_data(key.symbol, key.timeframe, days_back=self._days_back)
        except Exception:
            # Best-effort warmup; ignore errors so streaming can proceed
            pass

    def _resolve_partitions(self, key: ReplayKey) -> List[Tuple[datetime, Path]]:
        base_path = self._root / key.symbol / key.timeframe
        partitions: List[Tuple[datetime, Path]] = []
        if not base_path.exists():
            return partitions

        for part in base_path.glob("date=*"):
            if not part.is_dir():
                continue
            _, _, date_part = part.name.partition("=")
            try:
                dt = datetime.strptime(date_part, "%Y-%m-%d")
            except ValueError:
                continue
            partitions.append((dt, part))
        partitions.sort(key=lambda x: x[0])
        return partitions

    def _load_bars(self, key: ReplayKey) -> List[Dict[str, Any]]:
        cached = self._bar_cache.get(key)
        now = time.time()
        if cached and now - cached[0] < 300:
            return cached[1]

        partitions = self._resolve_partitions(key)
        if not partitions:
            self._bar_cache[key] = (now, [])
            return []

        cutoff_date = partitions[-1][0] - timedelta(days=self._days_back - 1)
        eligible = [p for p in partitions if p[0] >= cutoff_date]
        if not eligible:
            eligible = partitions[-self._days_back :] if len(partitions) >= self._days_back else partitions

        rows: List[Dict[str, Any]] = []
        for _, part in eligible:
            files = sorted(part.glob("*.parquet"))
            for file in files:
                try:
                    table = pq.read_table(file, columns=["time", "open", "high", "low", "close", "volume"])
                except Exception:
                    continue
                for row in table.to_pylist():
                    try:
                        rows.append(
                            {
                                "time": int(row["time"]),
                                "open": float(row["open"]),
                                "high": float(row["high"]),
                                "low": float(row["low"]),
                                "close": float(row["close"]),
                                "volume": float(row.get("volume") or 0.0),
                            }
                        )
                    except Exception:
                        continue

        rows.sort(key=lambda r: r["time"])
        self._bar_cache[key] = (now, rows)
        return rows

    async def subscribe(self, symbol: str, timeframe: str = "1m") -> Tuple[asyncio.Queue, Any]:
        if timeframe != "1m":
            raise ValueError("HistoricalReplayProvider currently supports timeframe '1m' only")
        key = ReplayKey(symbol=symbol.upper().replace("/", ""), timeframe=timeframe)
        queue: asyncio.Queue = asyncio.Queue(maxsize=200)
        live_store.subscribe(key.symbol, key.timeframe, queue)
        async with self._lock:
            self._active_counts[key] = self._active_counts.get(key, 0) + 1
            if key not in self._tasks:
                self._tasks[key] = asyncio.create_task(self._run_replay(key))

        async def _unsubscribe() -> None:
            live_store.unsubscribe(key.symbol, key.timeframe, queue)
            async with self._lock:
                self._active_counts[key] = max(0, self._active_counts.get(key, 0) - 1)
                if self._active_counts.get(key, 0) == 0 and key not in self._pinned:
                    task = self._tasks.pop(key, None)
                    if task:
                        task.cancel()
                    self._active_counts.pop(key, None)
        return queue, _unsubscribe

    async def ensure_stream(self, symbol: str, timeframe: str = "1m") -> None:
        if timeframe != "1m":
            raise ValueError("HistoricalReplayProvider currently supports timeframe '1m' only")
        key = ReplayKey(symbol=symbol.upper().replace("/", ""), timeframe=timeframe)
        async with self._lock:
            self._pinned.add(key)
            self._active_counts.setdefault(key, 0)
            if key not in self._tasks:
                self._tasks[key] = asyncio.create_task(self._run_replay(key))

    async def release_stream(self, symbol: str, timeframe: str = "1m") -> None:
        key = ReplayKey(symbol=symbol.upper().replace("/", ""), timeframe=timeframe)
        async with self._lock:
            self._pinned.discard(key)
            if self._active_counts.get(key, 0) == 0:
                task = self._tasks.pop(key, None)
                if task:
                    task.cancel()
                self._active_counts.pop(key, None)

    async def _run_replay(self, key: ReplayKey) -> None:
        self._ensure_history(key)
        sleep_secs = self._sleep_seconds(key.timeframe)
        tf_seconds = _TIMEFRAME_SECONDS.get(key.timeframe.lower(), 60)
        self._positions.setdefault(key, 0)
        self._last_times.setdefault(key, 0)

        while True:
            try:
                bars = self._load_bars(key)
                if not bars:
                    await asyncio.sleep(5.0)
                    continue

                idx = self._positions.get(key, 0)
                if idx >= len(bars):
                    if not self._loop:
                        await asyncio.sleep(sleep_secs)
                        continue
                    idx = 0
                bar = bars[idx]
                self._positions[key] = idx + 1

                now = int(time.time())
                aligned = (now // tf_seconds) * tf_seconds
                last_time = self._last_times.get(key) or 0
                if aligned <= last_time:
                    aligned = last_time + tf_seconds

                payload = {
                    "time": aligned,
                    "open": bar["open"],
                    "high": bar["high"],
                    "low": bar["low"],
                    "close": bar["close"],
                    "volume": bar.get("volume", 0.0),
                }
                live_store.ingest_bar(symbol=key.symbol, timeframe=key.timeframe, bar=payload)
                self._last_times[key] = aligned

                await asyncio.sleep(sleep_secs)
            except asyncio.CancelledError:
                break
            except Exception:
                await asyncio.sleep(min(5.0, sleep_secs))


historical_replay = HistoricalReplayProvider()
