from __future__ import annotations

import asyncio
import json
from typing import List
import contextlib

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query

from app.streaming.binance_hub import hub as binance_hub
from app.streaming.alpha_hub import alpha_hub

router = APIRouter(prefix="/stream", tags=["Stream"])


@router.websocket("/klines")
async def stream_klines_ws(
    websocket: WebSocket,
    symbols: str = Query("BTCUSDT,XAUUSD"),
    timeframe: str = Query("1m"),
):
    await websocket.accept()

    syms: List[str] = [s.strip().upper().replace("/", "") for s in symbols.split(",") if s.strip()]


    async def forward_alpha(sym: str):
        queue, unsubscribe = await alpha_hub.subscribe(sym, timeframe="1m")
        try:
            while True:
                evt = await queue.get()
                if evt is None:
                    break
                try:
                    await websocket.send_text(json.dumps(evt))
                except WebSocketDisconnect:
                    break
                except Exception:
                    await asyncio.sleep(0)
        finally:
            try:
                await unsubscribe()
            except Exception:
                pass

    async def forward_binance(sym: str):
        queue, unsubscribe = await binance_hub.subscribe(sym, timeframe="1m")
        try:
            while True:
                evt = await queue.get()
                if evt is None:
                    break
                # evt already normalized with iso timestamps
                try:
                    await websocket.send_text(json.dumps(evt))
                except WebSocketDisconnect:
                    break
                except Exception:
                    await asyncio.sleep(0)
        finally:
            try:
                await unsubscribe()
            except Exception:
                pass

    tasks: List[asyncio.Task] = []
    try:
        for s in syms:
            if s == "XAUUSD" or s.startswith("XAU"):
                tasks.append(asyncio.create_task(forward_alpha("XAUUSD")))
            else:
                tasks.append(asyncio.create_task(forward_binance(s)))

        # Wait for disconnect
        done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION)
    except WebSocketDisconnect:
        pass
    finally:
        for t in tasks:
            t.cancel()
            with contextlib.suppress(Exception):
                await t
        with contextlib.suppress(Exception):
            await websocket.close()