from __future__ import annotations

import asyncio
import json
import os
from dataclasses import dataclass
from datetime import datetime
from typing import Dict, Set, Tuple, Any

import websockets

from app.streaming.live_store import live_store


@dataclass(frozen=True)
class StreamKey:
    symbol: str
    timeframe: str  # only '1m' supported in hub


class BinanceStreamHub:
    """
    Maintains a single upstream websocket per (symbol,timeframe) and fans out
    kline events to multiple subscribers via asyncio.Queues.
    """

    def __init__(self, base_ws: str | None = None) -> None:
        self.base_ws = (base_ws or os.getenv("BINANCE_WS_URL", "wss://stream.binance.com:9443/ws")).rstrip("/")
        self._subs: Dict[StreamKey, Set[asyncio.Queue]] = {}
        self._tasks: Dict[StreamKey, asyncio.Task] = {}
        self._lock = asyncio.Lock()

    def get_status(self) -> list[dict]:
        """Return status snapshot of active streams."""
        out: list[dict] = []
        for key, subs in self._subs.items():
            hist = live_store.get_history(key.symbol, key.timeframe)
            last_ts = hist[-1]["time"] if hist else None
            last_iso = None
            if isinstance(last_ts, (int, float)):
                try:
                    last_iso = datetime.utcfromtimestamp(int(last_ts)).isoformat() + "Z"
                except Exception:
                    last_iso = None
            out.append({
                "symbol": key.symbol,
                "timeframe": key.timeframe,
                "subscribers": len(subs),
                "last_event_time": last_iso,
            })
        return out

    async def subscribe(self, symbol: str, timeframe: str = "1m") -> Tuple[asyncio.Queue, Any]:
        """Subscribe to a stream. Returns (queue, unsubscribe_cb)."""
        if timeframe != "1m":
            raise ValueError("BinanceStreamHub currently supports timeframe '1m' only")
        key = StreamKey(symbol=symbol.upper().replace("/", ""), timeframe=timeframe)
        q: asyncio.Queue = asyncio.Queue(maxsize=1000)
        async with self._lock:
            subs = self._subs.get(key)
            if not subs:
                subs = set()
                self._subs[key] = subs
            subs.add(q)
            if key not in self._tasks:
                self._tasks[key] = asyncio.create_task(self._run_stream(key))

        async def _unsubscribe() -> None:
            async with self._lock:
                s = self._subs.get(key)
                if s and q in s:
                    s.remove(q)
                # Close queue to unblock listeners
                try:
                    q.put_nowait(None)
                except Exception:
                    pass
                if s is not None and len(s) == 0:
                    # cancel task and cleanup
                    t = self._tasks.pop(key, None)
                    if t:
                        t.cancel()
                    self._subs.pop(key, None)
        return q, _unsubscribe

    async def _run_stream(self, key: StreamKey) -> None:
        symbol = key.symbol
        stream = f"{symbol.lower()}@kline_{key.timeframe}"
        url = self.base_ws.replace("/ws", "/stream") + f"?streams={stream}"
        # Reconnect loop
        while True:
            try:
                async with websockets.connect(url, ping_interval=20, ping_timeout=20) as ws:
                    async for message in ws:
                        try:
                            data = json.loads(message)
                            k = (data.get("data") or {}).get("k") or {}
                            if not k:
                                continue
                            # Normalize event
                            evt = {
                                "symbol": symbol,
                                "timeframe": key.timeframe,
                                "open_time": datetime.fromtimestamp(k["t"] / 1000.0).isoformat(),
                                "close_time": datetime.fromtimestamp(k["T"] / 1000.0).isoformat(),
                                "open": float(k["o"]),
                                "high": float(k["h"]),
                                "low": float(k["l"]),
                                "close": float(k["c"]),
                                "volume": float(k.get("v", 0.0)),
                                "is_closed": bool(k.get("x", False)),
                                "source": "binance",
                            }
                            # Update live store (1m bar)
                            try:
                                tsec = int(k["T"] // 1000)
                                live_store.ingest_bar(symbol=symbol, timeframe=key.timeframe, bar={
                                    "time": tsec, "open": evt["open"], "high": evt["high"], "low": evt["low"], "close": evt["close"], "volume": evt["volume"],
                                })
                            except Exception:
                                pass
                            # Fan-out to subscribers
                            subs = self._subs.get(key) or set()
                            for q in list(subs):
                                try:
                                    if q.full():
                                        q.get_nowait()
                                    q.put_nowait(evt)
                                except Exception:
                                    # Drop failed subscriber
                                    try:
                                        subs.remove(q)
                                    except Exception:
                                        pass
                        except Exception:
                            continue
            except Exception:
                await asyncio.sleep(1.5)


# Singleton hub instance
hub = BinanceStreamHub()