chart/app/api/ws.py

import asyncio
from urllib.parse import urlsplit

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from app.api.deps import SESSION_COOKIE, session_matches, token_matches
from app.analysis.levels import LevelKind
from app.bars.models import Timeframe
from app.analysis.bar_space import fill_short_gaps
from app.bars.session import bucket_duration, future_bucket_starts
from app.analysis.confluence import cluster_levels

router = APIRouter()


def same_origin(websocket: WebSocket) -> bool:
    origin = websocket.headers.get("origin", "")
    host = websocket.headers.get("host", "")
    return bool(origin and host) and urlsplit(origin).netloc == host


def level_enabled(level, enabled: dict) -> bool:
    kind = level.kind.value
    if kind == "ma":
        return level.period in enabled.get("ma", {}).get(level.tf.value, [])
    if kind == "manual":
        return enabled.get("drawings", True) and enabled.get("manual", True)
    if kind == "trendline":
        return enabled.get("auto", False)
    if kind == "horizontal":
        return enabled.get("horizontal", True)
    if kind == "vwap":
        return enabled.get("vwap", True)
    return False


def enabled_levels(runtime, prefs: dict | None):
    if not prefs or prefs.get("hidden_levels_score"):
        return runtime.levels
    enabled = prefs.get("enabled", {})
    return [level for level in runtime.levels if level_enabled(level, enabled)]


def connection_clusters(runtime, prefs: dict | None):
    if runtime.price is None or runtime.stream.last_bar_t is None:
        return []
    return cluster_levels(
        enabled_levels(runtime, prefs), runtime.stream.last_bar_t, runtime.price, runtime.atr15
    )


def session_open(runtime) -> float | None:
    bars = runtime.store.get(Timeframe.D1, 1)
    return bars[-1].o if bars else None


def session_range(runtime) -> dict | None:
    """High/low of the forming CME session — today's range, not prior day."""
    bars = runtime.store.get(Timeframe.D1, 1)
    if not bars:
        return None
    bar = bars[-1]
    return {"t": bar.t, "high": bar.h, "low": bar.l}


def trendline_timeframes(runtime) -> set[Timeframe]:
    return {
        level.tf for level in runtime.levels
        if level.kind is LevelKind.MANUAL and level.slope
    }


def trendline_geometry(runtime) -> dict:
    if not runtime.settings.trendline_source_geometry:
        return {"mode": "legacy", "series": {}}
    series = {}
    for tf in sorted(trendline_timeframes(runtime), key=lambda value: value.value):
        series[tf.value] = trendline_series(runtime, tf)
    return {"mode": "source_tf", "series": series}


def trendline_series(runtime, tf: Timeframe) -> dict:
    times = fill_short_gaps([bar.t for bar in runtime.store.get(tf)], tf)
    value = {"times": times}
    if tf is Timeframe.D1:
        value["durations"] = [bucket_duration(t, tf) for t in times]
    else:
        value["duration"] = tf.seconds
    future = future_bucket_starts(times[-1], tf) if times else []
    value["future_times"] = future
    value["future_durations"] = [bucket_duration(t, tf) for t in future]
    return value


def displayed_future_times(runtime, tf: Timeframe) -> list[int]:
    bars = runtime.store.get(tf, 1)
    return future_bucket_starts(bars[-1].t, tf) if bars else []


def bar_client_message(runtime, bar, displayed_tf: Timeframe, *, full: bool) -> dict | None:
    """What one socket should hear about a bar.

    A forming tick of an already-seen minute is just the live candle. Future
    calendars and HTF geometry go out only when that timeframe's timestamp
    advances — otherwise 1m at 4 Hz resends 180 slots and walks the store.
    """
    if bar.tf is displayed_tf:
        payload = {"type": "bar", "tf": displayed_tf.value, "bar": bar.to_dict()}
        if not full:
            return payload
        source_tfs = trendline_timeframes(runtime)
        source_bars = runtime.store.get(bar.tf)
        source_geometry = trendline_series(runtime, bar.tf) if source_bars else {}
        first_source_t = source_bars[0].t if source_bars else None
        source_index = next(
            (index for index, value in enumerate(source_bars) if value.t == bar.t), None,
        )
        previous_source_t = (
            source_bars[source_index - 1].t
            if source_index is not None and source_index > 0 else None
        )
        rng = session_range(runtime)
        extra = {
            "session_t": rng["t"],
            "session_high": rng["high"],
            "session_low": rng["low"],
        } if rng else {}
        payload.update(
            {
                "duration": bucket_duration(bar.t, bar.tf),
                "session_open": session_open(runtime),
                **extra,
                "trendline_first_t": (
                    first_source_t
                    if runtime.settings.trendline_source_geometry
                    and bar.tf in source_tfs else None
                ),
                "trendline_previous_t": previous_source_t,
                "trendline_future_times": source_geometry.get("future_times", []),
                "trendline_future_durations": source_geometry.get("future_durations", []),
                "future_times": source_geometry.get("future_times", []),
            }
        )
        return payload
    if (
        full
        and runtime.settings.trendline_source_geometry
        and bar.tf in trendline_timeframes(runtime)
    ):
        source_bars = runtime.store.get(bar.tf)
        source_geometry = trendline_series(runtime, bar.tf) if source_bars else {}
        first_source_t = source_bars[0].t if source_bars else None
        source_index = next(
            (index for index, value in enumerate(source_bars) if value.t == bar.t), None,
        )
        previous_source_t = (
            source_bars[source_index - 1].t
            if source_index is not None and source_index > 0 else None
        )
        return {
            "type": "trendline_bar",
            "tf": bar.tf.value,
            "t": bar.t,
            "duration": bucket_duration(bar.t, bar.tf),
            "first_t": first_source_t,
            "previous_t": previous_source_t,
            "future_times": source_geometry.get("future_times", []),
            "future_durations": source_geometry.get("future_durations", []),
        }
    return None


def snapshot(runtime, tf: Timeframe, prefs: dict | None = None) -> dict:
    events, events_more = runtime.events.recent()
    rng = session_range(runtime)
    extra = {
        "session_t": rng["t"],
        "session_high": rng["high"],
        "session_low": rng["low"],
    } if rng else {}
    return {
        "type": "snapshot",
        "tf": tf.value,
        "bars": [bar.to_dict() for bar in runtime.store.get(tf)],
        "levels": [level.to_dict() for level in runtime.levels],
        "clusters": [cluster.to_dict() for cluster in connection_clusters(runtime, prefs)],
        "price": runtime.store.get(Timeframe.M1, 1)[-1].c
        if runtime.store.get(Timeframe.M1, 1)
        else None,
        "session_open": session_open(runtime),
        **extra,
        "trendline_geometry": trendline_geometry(runtime),
        "future_times": displayed_future_times(runtime, tf),
        "events": events,
        "events_more": events_more,
        "instrument": runtime.settings.profile.payload(),
    }


@router.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    # Browsers automatically include the HttpOnly session cookie in the
    # handshake. Query-token support remains for non-browser clients and for
    # tabs migrating from the previous localStorage-based login.
    token_ok = token_matches(websocket.app, websocket.query_params.get("token", ""))
    session_ok = same_origin(websocket) and session_matches(
        websocket.app, websocket.cookies.get(SESSION_COOKIE, "")
    )
    if not token_ok and not session_ok:
        await websocket.close(code=1008, reason="Missing or invalid chart token")
        return
    await websocket.accept()
    runtime = websocket.app.state.runtime
    queue: asyncio.Queue = asyncio.Queue(maxsize=100)
    runtime.subscribers.add(queue)
    tf = Timeframe.M1
    prefs = None
    await websocket.send_json(snapshot(runtime, tf, prefs))
    geometry_tfs = trendline_timeframes(runtime)
    last_bar_t: dict[Timeframe, int] = {}

    async def receive():
        nonlocal tf, prefs
        try:
            while True:
                message = await websocket.receive_json()
                if message.get("type") == "subscribe":
                    tf = Timeframe(message.get("tf", "1m"))
                    await websocket.send_json(snapshot(runtime, tf, prefs))
                elif message.get("type") == "prefs":
                    prefs = message
                    clusters = connection_clusters(runtime, prefs)
                    await websocket.send_json(
                        {
                            "type": "clusters",
                            "price": runtime.price,
                            "session_open": session_open(runtime),
                            "clusters": [cluster.to_dict() for cluster in clusters],
                        }
                    )
                elif message.get("type") == "trendline_geometry":
                    await websocket.send_json(
                        {"type": "trendline_geometry", "geometry": trendline_geometry(runtime)}
                    )
        except WebSocketDisconnect:
            queue.put_nowait({"type": "disconnect"})

    receiver = asyncio.create_task(receive())
    try:
        while True:
            event = await queue.get()
            if event["type"] == "disconnect":
                break
            if event["type"] == "resync":
                # History changed behind the live edge (a backfilled outage).
                # Bar deltas only move the tail, so the whole series is resent.
                last_bar_t.clear()
                await websocket.send_json(snapshot(runtime, tf, prefs))
            elif event["type"] == "bar":
                bar = event["bar"]
                full = last_bar_t.get(bar.tf) != bar.t
                if full:
                    last_bar_t[bar.tf] = bar.t
                payload = bar_client_message(runtime, bar, tf, full=full)
                if payload is not None:
                    await websocket.send_json(payload)
            elif event["type"] == "levels":
                message = {
                    "type": "levels", "changed": event["changed"], "removed": event["removed"],
                }
                current_tfs = trendline_timeframes(runtime)
                if current_tfs != geometry_tfs:
                    message["trendline_geometry"] = trendline_geometry(runtime)
                    geometry_tfs = current_tfs
                await websocket.send_json(message)
            elif event["type"] == "clusters":
                clusters = connection_clusters(runtime, prefs)
                await websocket.send_json(
                    {
                        "type": "clusters",
                        "price": runtime.price,
                        "session_open": session_open(runtime),
                        "clusters": [cluster.to_dict() for cluster in clusters],
                    }
                )
            elif event["type"] == "alert":
                # Alerts are produced once, server-side. This socket only relays
                # them, so opening a second tab cannot double-notify.
                await websocket.send_json(
                    {
                        "type": "alert",
                        "number": event.get("number", 0),
                        "at": event.get("at", 0),
                        "cluster": event["cluster"].to_dict(),
                        "message": event["message"],
                    }
                )
    except (WebSocketDisconnect, asyncio.CancelledError):
        pass
    finally:
        receiver.cancel()
        runtime.subscribers.discard(queue)