"""AngelOne SmartStream WebSocket client and binary tick parser."""

from __future__ import annotations

import asyncio
import json
import logging
import struct
from datetime import datetime
from typing import Any

import websockets
import websockets.exceptions

from tt_connect.core.timezone import IST
from tt_connect.core.store.resolver import ResolvedInstrument
from tt_connect.core.models.instruments import Instrument
from tt_connect.core.models import Tick
from tt_connect.core.models.enums import FeedState
from tt_connect.core.adapter.ws import BrokerWebSocket

logger = logging.getLogger(__name__)

_WS_URL = "wss://smartapisocket.angelone.in/smart-stream"

# Exchange → AngelOne exchangeType integer
_EXCHANGE_TYPE: dict[str, int] = {
    "NSE": 1,
    "NFO": 2,
    "BSE": 3,
    "BFO": 4,
    "MCX": 5,
}

# Subscription modes
_MODE_LTP        = 1
_MODE_QUOTE      = 2
_MODE_SNAP_QUOTE = 3

# Minimum packet sizes per mode (little-endian binary)
_LTP_MIN        = 51   # bytes 0-50
_QUOTE_MIN      = 123  # bytes 0-122
_SNAP_QUOTE_MIN = 347  # bytes 0-346 (need best_5 block to parse bid/ask)


class AngelOneWebSocket(BrokerWebSocket):
    """
    AngelOne SmartStream WebSocket client.

    Binary protocol (little-endian):
      byte  0:     subscription_mode (B)
      byte  1:     exchange_type (B)
      bytes 2-26:  token (25-byte null-terminated ASCII)
      bytes 27-34: sequence_number (q)
      bytes 35-42: exchange_timestamp (q, milliseconds)
      bytes 43-50: last_traded_price (q, paise → ÷100)

      QUOTE (mode≥2, ≥123 bytes) adds:
        bytes 51-58:   last_traded_quantity (q)
        bytes 59-66:   average_traded_price (q, paise)
        bytes 67-74:   volume_trade_for_the_day (q)
        bytes 75-82:   total_buy_quantity (d)
        bytes 83-90:   total_sell_quantity (d)
        bytes 91-98:   open_price (q, paise)
        bytes 99-106:  high_price (q, paise)
        bytes 107-114: low_price (q, paise)
        bytes 115-122: closed_price (q, paise)

      SNAP_QUOTE (mode=3, ≥347 bytes) adds:
        bytes 123-130: last_traded_timestamp (q)
        bytes 131-138: open_interest (q)
        bytes 139-146: oi_change_percentage (q)
        bytes 147-346: best_5_buy_sell (200 bytes, 10×20)
                       each 20-byte packet: flag(H) qty(q) price(q) orders(H)
                       flag==0 → buy side, flag!=0 → sell side
    """

    PING_INTERVAL       = 10    # seconds between heartbeat text pings
    STALE_THRESHOLD     = 30    # seconds without a tick before feed → STALE
    MAX_RECONNECT_DELAY = 60    # seconds
    _BROKER_NAME        = "angelone"

    def __init__(self, auth: Any) -> None:
        super().__init__()
        # auth: AngelOneAuth (typed loosely to avoid circular import)
        self._auth = auth

        # token (str) → Instrument  — reverse map for incoming ticks
        self._token_map: dict[str, Instrument] = {}
        # token (str) → exchangeType (int) — needed to build subscribe messages
        self._token_exchange_type: dict[str, int] = {}

    # ------------------------------------------------------------------ abstract hook implementations

    async def _register_subscriptions(
        self, subscriptions: list[tuple[Instrument, ResolvedInstrument]]
    ) -> list[Any]:
        """Store token→instrument mappings; return new token strings."""
        new_tokens: list[str] = []
        for instrument, resolved in subscriptions:
            self._token_map[resolved.token] = instrument
            self._token_exchange_type[resolved.token] = _EXCHANGE_TYPE.get(resolved.exchange, 1)
            new_tokens.append(resolved.token)
        return new_tokens

    def _tokens_for_instruments(self, instruments: list[Instrument]) -> list[Any]:
        return [t for t, inst in self._token_map.items() if inst in instruments]

    def _remove_tokens(self, tokens: list[Any]) -> None:
        for t in tokens:
            self._token_map.pop(t, None)
            self._token_exchange_type.pop(t, None)

    def _all_tracked_tokens(self) -> list[Any]:
        return list(self._token_map.keys())

    async def _maybe_ping(self, ws: Any) -> None:
        """Send AngelOne application-level text ping."""
        await ws.send("ping")

    # ------------------------------------------------------------ main loop

    async def _connect_and_run(self) -> None:
        """Open websocket, resubscribe tracked tokens, and dispatch ticks."""
        if not self._auth._session:
            raise RuntimeError("AngelOne auth session not available — call login() first")

        headers = {
            "Authorization":  self._auth.access_token,
            "x-api-key":      self._auth._config.get("api_key", ""),
            "x-client-code":  self._auth._config.get("client_id", ""),
            "x-feed-token":   self._auth._session.feed_token or "",
        }

        async with websockets.connect(
            _WS_URL,
            additional_headers=headers,
            ping_interval=None,  # AngelOne uses text "ping" frames — disable library ping/pong
        ) as ws:
            self._ws = ws
            self._feed_state = FeedState.CONNECTED
            logger.info("AngelOne WS connected", extra={"event": "ws.connect", "broker": "angelone"})

            # On (re)connect, resubscribe to everything currently tracked
            if self._token_map:
                await self._send_subscribe(ws, self._all_tracked_tokens())

            staleness_task: asyncio.Task[None] = asyncio.create_task(self._staleness_loop(ws))
            try:
                async for message in ws:
                    if isinstance(message, bytes):
                        tick = self._parse_binary(message)
                        if tick:
                            was_stale = self._record_tick(tick.instrument)
                            await self._maybe_fire_recovered(was_stale)
                            if self._on_tick:
                                try:
                                    await self._on_tick(tick)
                                except Exception:
                                    logger.exception(
                                        "AngelOne WS on_tick callback failed",
                                        extra={"event": "ws.tick_error", "broker": "angelone"},
                                    )
                    # text "pong" — ignore
            finally:
                staleness_task.cancel()
                try:
                    await staleness_task
                except asyncio.CancelledError:
                    pass
                self._ws = None
                logger.info("AngelOne WS disconnected", extra={"event": "ws.disconnect", "broker": "angelone"})

    # ------------------------------------------------- subscribe / unsubscribe

    async def _send_subscribe(
        self, ws: Any, tokens: list[Any], mode: int = _MODE_SNAP_QUOTE
    ) -> None:
        token_list = self._build_token_list(tokens)
        if not token_list:
            return
        await ws.send(json.dumps({
            "correlationID": "ttconnect",
            "action": 1,
            "params": {"mode": mode, "tokenList": token_list},
        }))
        logger.debug(
            f"Subscribed to {len(tokens)} tokens (mode={mode})",
            extra={"event": "ws.subscribe", "broker": "angelone", "token_count": len(tokens), "mode": mode},
        )

    async def _send_unsubscribe(
        self, ws: Any, tokens: list[Any], mode: int = _MODE_QUOTE
    ) -> None:
        token_list = self._build_token_list(tokens)
        if not token_list:
            return
        await ws.send(json.dumps({
            "correlationID": "ttconnect",
            "action": 0,
            "params": {"mode": mode, "tokenList": token_list},
        }))
        logger.debug(
            f"Unsubscribed from {len(tokens)} tokens",
            extra={"event": "ws.unsubscribe", "broker": "angelone", "token_count": len(tokens)},
        )

    def _build_token_list(self, tokens: list[Any]) -> list[dict[str, Any]]:
        """Group tokens by exchangeType for the AngelOne subscribe payload."""
        by_exchange: dict[int, list[str]] = {}
        for token in tokens:
            et = self._token_exchange_type.get(token, 1)
            by_exchange.setdefault(et, []).append(token)
        return [
            {"exchangeType": et, "tokens": toks}
            for et, toks in by_exchange.items()
        ]

    # ------------------------------------------------------ binary parser

    def _parse_binary(self, data: bytes) -> Tick | None:
        """Parse one binary SmartStream packet into canonical Tick."""
        if len(data) < _LTP_MIN:
            return None

        subscription_mode = data[0]
        token = data[2:27].split(b"\x00")[0].decode("ascii", errors="replace").strip()

        instrument = self._token_map.get(token)
        if instrument is None:
            logger.debug(f"Tick for unknown token {token!r}")
            return None

        ts_ms  = struct.unpack_from("<q", data, 35)[0]
        ltp    = struct.unpack_from("<q", data, 43)[0] / 100.0
        ts     = (
            datetime.fromtimestamp(ts_ms / 1000, tz=IST)
            if ts_ms > 0 else None
        )

        volume: int | None = None
        oi: int | None     = None
        bid: float | None  = None
        ask: float | None  = None

        if subscription_mode >= _MODE_QUOTE and len(data) >= _QUOTE_MIN:
            volume = struct.unpack_from("<q", data, 67)[0]

        if subscription_mode >= _MODE_SNAP_QUOTE and len(data) >= _SNAP_QUOTE_MIN:
            oi = struct.unpack_from("<q", data, 131)[0]
            bid, ask = self._parse_best5_top(data[147:347])

        return Tick(
            instrument=instrument,
            ltp=ltp,
            volume=volume,
            oi=oi,
            bid=bid,
            ask=ask,
            timestamp=ts,
        )

    @staticmethod
    def _parse_best5_top(block: bytes) -> tuple[float | None, float | None]:
        """
        Parse the 200-byte best-5 block and return (best_bid, best_ask).

        Each 20-byte packet:  flag(H=2) | quantity(q=8) | price(q=8) | orders(H=2)
        flag == 0  → buy side (we want the first one = highest bid)
        flag != 0  → sell side (we want the first one = lowest ask)
        Prices are in paise; divide by 100 for rupees.
        """
        bid: float | None = None
        ask: float | None = None
        for i in range(10):
            offset = i * 20
            if offset + 20 > len(block):
                break
            flag, _qty, price_paise, _orders = struct.unpack_from("<HqqH", block, offset)
            price = price_paise / 100.0
            if price <= 0:
                continue
            if flag == 0 and bid is None:
                bid = price
            elif flag != 0 and ask is None:
                ask = price
            if bid is not None and ask is not None:
                break
        return bid, ask
