from __future__ import annotations

import threading
import time
from typing import Iterable

from .types import BookLevel, MarketSnapshot, OptionPair, Quote


class BookCache:
    """内存盘口缓存：永续 + Call/Put。线程安全。"""

    def __init__(self) -> None:
        self._lock = threading.RLock()
        self._quotes: dict[str, Quote] = {}
        self._index_px: float | None = None
        self._pair: OptionPair | None = None
        self._connected = False
        self._updated_at_ms: int | None = None

    def set_connected(self, ok: bool) -> None:
        with self._lock:
            self._connected = bool(ok)

    def set_pair(self, pair: OptionPair | None) -> None:
        with self._lock:
            self._pair = pair

    def set_index_px(self, px: float | None) -> None:
        with self._lock:
            if px is not None and px > 0:
                self._index_px = float(px)
                self._touch()

    def upsert_book(
        self,
        inst_id: str,
        *,
        bids: list[BookLevel],
        asks: list[BookLevel],
        ts_ms: int | None = None,
    ) -> None:
        with self._lock:
            q = self._quotes.get(inst_id) or Quote(inst_id=inst_id)
            q.bids = bids
            q.asks = asks
            q.bid = bids[0].px if bids else None
            q.ask = asks[0].px if asks else None
            q.bid_sz = bids[0].sz if bids else None
            q.ask_sz = asks[0].sz if asks else None
            if ts_ms is not None:
                q.ts_ms = ts_ms
            self._quotes[inst_id] = q
            self._touch(ts_ms)

    def upsert_top(
        self,
        inst_id: str,
        *,
        bid: float | None,
        ask: float | None,
        bid_sz: float | None = None,
        ask_sz: float | None = None,
        ts_ms: int | None = None,
    ) -> None:
        with self._lock:
            q = self._quotes.get(inst_id) or Quote(inst_id=inst_id)
            if bid is not None:
                q.bid = bid
            if ask is not None:
                q.ask = ask
            if bid_sz is not None:
                q.bid_sz = bid_sz
            if ask_sz is not None:
                q.ask_sz = ask_sz
            if ts_ms is not None:
                q.ts_ms = ts_ms
            # 同步一层盘口，便于 snapshot 展示
            if bid is not None and bid_sz is not None:
                q.bids = [BookLevel(px=bid, sz=bid_sz)] + q.bids[1:]
            if ask is not None and ask_sz is not None:
                q.asks = [BookLevel(px=ask, sz=ask_sz)] + q.asks[1:]
            self._quotes[inst_id] = q
            self._touch(ts_ms)

    def set_mark_px(self, inst_id: str, mark_px: float | None, ts_ms: int | None = None) -> None:
        with self._lock:
            if mark_px is None or mark_px <= 0:
                return
            q = self._quotes.get(inst_id) or Quote(inst_id=inst_id)
            q.mark_px = float(mark_px)
            if ts_ms is not None:
                q.ts_ms = ts_ms
            self._quotes[inst_id] = q
            self._touch(ts_ms)

    def get(self, inst_id: str) -> Quote | None:
        with self._lock:
            return self._quotes.get(inst_id)

    def drop_except(self, keep: Iterable[str]) -> None:
        keep_set = set(keep)
        with self._lock:
            for k in list(self._quotes):
                if k not in keep_set:
                    del self._quotes[k]

    def snapshot(self, perp_inst_id: str) -> MarketSnapshot:
        with self._lock:
            pair = self._pair
            call = self._quotes.get(pair.call_inst_id) if pair else None
            put = self._quotes.get(pair.put_inst_id) if pair else None
            return MarketSnapshot(
                perp=self._quotes.get(perp_inst_id),
                call=call,
                put=put,
                index_px=self._index_px,
                pair=pair,
                connected=self._connected,
                updated_at_ms=self._updated_at_ms,
            )

    def _touch(self, ts_ms: int | None = None) -> None:
        self._updated_at_ms = int(ts_ms) if ts_ms is not None else int(time.time() * 1000)