"""In-memory L2 order book per token.

Kept deliberately small: a sorted dict per side (price -> size), plus a bit of
bookkeeping to detect desync. The opportunity engine consumes immutable
`Book` pydantic snapshots rendered from this structure on every update.
"""
from __future__ import annotations

import asyncio
from collections.abc import AsyncIterator
from dataclasses import dataclass, field
from datetime import UTC, datetime
from decimal import Decimal
from enum import Enum
from typing import Self

from sortedcontainers import SortedDict

from ..models import Book, BookLevel

REMOVE_AT_ZERO = Decimal(0)


class Side(str, Enum):
    BID = "BID"
    ASK = "ASK"


def _to_decimal(v: object) -> Decimal:
    if isinstance(v, Decimal):
        return v
    if isinstance(v, (int, float)):
        return Decimal(str(v))
    if isinstance(v, str):
        return Decimal(v)
    raise TypeError(f"cannot coerce {type(v).__name__} to Decimal")


def _coerce_side(v: object) -> Side:
    if isinstance(v, Side):
        return v
    s = str(v).strip().upper()
    if s in ("BUY", "BID", "B"):
        return Side.BID
    if s in ("SELL", "ASK", "S", "A"):
        return Side.ASK
    raise ValueError(f"unknown side: {v!r}")


@dataclass(slots=True)
class LevelChange:
    price: Decimal
    size: Decimal
    side: Side

    @classmethod
    def from_raw(cls, raw: dict[str, object]) -> Self:
        return cls(
            price=_to_decimal(raw["price"]),
            size=_to_decimal(raw["size"]),
            side=_coerce_side(raw["side"]),
        )


@dataclass(slots=True)
class LiveBook:
    """Mutable L2 book for one token. Bids and asks are sorted price ladders."""

    token_id: str
    bids: SortedDict = field(default_factory=SortedDict)  # price (Decimal) -> size (Decimal)
    asks: SortedDict = field(default_factory=SortedDict)
    snapshots_applied: int = 0
    deltas_applied: int = 0
    last_update: datetime | None = None
    last_hash: str | None = None

    def apply_snapshot(
        self,
        *,
        bids: list[tuple[Decimal, Decimal]],
        asks: list[tuple[Decimal, Decimal]],
        timestamp: datetime | None = None,
        book_hash: str | None = None,
    ) -> None:
        self.bids.clear()
        self.asks.clear()
        for price, size in bids:
            if size > 0:
                self.bids[price] = size
        for price, size in asks:
            if size > 0:
                self.asks[price] = size
        self.snapshots_applied += 1
        self.last_update = timestamp or datetime.now(UTC)
        self.last_hash = book_hash

    def apply_delta(
        self,
        changes: list[LevelChange],
        *,
        timestamp: datetime | None = None,
        book_hash: str | None = None,
    ) -> None:
        for ch in changes:
            side_map = self.bids if ch.side is Side.BID else self.asks
            if ch.size <= REMOVE_AT_ZERO:
                side_map.pop(ch.price, None)
            else:
                side_map[ch.price] = ch.size
        self.deltas_applied += 1
        self.last_update = timestamp or datetime.now(UTC)
        self.last_hash = book_hash

    def best_bid(self) -> tuple[Decimal, Decimal] | None:
        if not self.bids:
            return None
        price = self.bids.keys()[-1]
        return (price, self.bids[price])

    def best_ask(self) -> tuple[Decimal, Decimal] | None:
        if not self.asks:
            return None
        price = self.asks.keys()[0]
        return (price, self.asks[price])

    def vwap_buy(self, target_size: Decimal) -> tuple[Decimal, Decimal, int] | None:
        """Walk asks ascending to fill `target_size`.
        Returns (vwap_price, filled_size, levels_consumed). If unable to fill any, None.
        """
        if target_size <= 0 or not self.asks:
            return None
        remaining = target_size
        cost = Decimal(0)
        filled = Decimal(0)
        levels = 0
        for price in self.asks.keys():
            size = self.asks[price]
            take = min(remaining, size)
            cost += take * price
            filled += take
            levels += 1
            remaining -= take
            if remaining <= 0:
                break
        if filled <= 0:
            return None
        return (cost / filled, filled, levels)

    def to_snapshot(self) -> Book:
        bids = [
            BookLevel(price=p, size=s)
            for p, s in reversed(list(self.bids.items()))  # descending
        ]
        asks = [BookLevel(price=p, size=s) for p, s in self.asks.items()]
        return Book(
            token_id=self.token_id,
            bids=bids,
            asks=asks,
            sequence=self.snapshots_applied + self.deltas_applied,
            updated_at=self.last_update or datetime.now(UTC),
        )


@dataclass(slots=True)
class BookUpdate:
    """Emitted whenever a token's book is mutated."""

    token_id: str
    reason: str  # "snapshot" | "delta"
    at: datetime


class BookRegistry:
    """Holds LiveBooks and multiplexes update events to subscribers."""

    def __init__(self) -> None:
        self._books: dict[str, LiveBook] = {}
        self._subscribers: list[asyncio.Queue[BookUpdate]] = []

    def book(self, token_id: str) -> LiveBook:
        book = self._books.get(token_id)
        if book is None:
            book = LiveBook(token_id=token_id)
            self._books[token_id] = book
        return book

    def get(self, token_id: str) -> LiveBook | None:
        return self._books.get(token_id)

    def tokens(self) -> list[str]:
        return list(self._books.keys())

    def apply_snapshot(
        self,
        token_id: str,
        *,
        bids: list[tuple[Decimal, Decimal]],
        asks: list[tuple[Decimal, Decimal]],
        timestamp: datetime | None = None,
        book_hash: str | None = None,
    ) -> None:
        self.book(token_id).apply_snapshot(
            bids=bids, asks=asks, timestamp=timestamp, book_hash=book_hash
        )
        self._publish(BookUpdate(token_id=token_id, reason="snapshot", at=datetime.now(UTC)))

    def apply_delta(
        self,
        token_id: str,
        changes: list[LevelChange],
        *,
        timestamp: datetime | None = None,
        book_hash: str | None = None,
    ) -> None:
        self.book(token_id).apply_delta(
            changes, timestamp=timestamp, book_hash=book_hash
        )
        self._publish(BookUpdate(token_id=token_id, reason="delta", at=datetime.now(UTC)))

    async def updates(self, maxsize: int = 1024) -> AsyncIterator[BookUpdate]:
        queue: asyncio.Queue[BookUpdate] = asyncio.Queue(maxsize=maxsize)
        self._subscribers.append(queue)
        try:
            while True:
                yield await queue.get()
        finally:
            self._subscribers.remove(queue)

    def _publish(self, update: BookUpdate) -> None:
        for q in self._subscribers:
            try:
                q.put_nowait(update)
            except asyncio.QueueFull:
                # If a consumer can't keep up, drop — the engine uses the
                # *current* book state anyway, so missing a tick is harmless.
                pass