#!/usr/bin/env python3
"""
Market Data WebSocket Server — simulates real-time stock price feeds.

Features:
  - Subscribe/unsubscribe per symbol
  - Periodic price updates with random walk
  - Timestamped messages for client-side latency measurement
  - Push interval: configurable (default 100ms = 10 msg/sec per symbol)

Usage:
  python3 market_data_server.py [port] [push_interval_ms]

Clients send:  {"action":"subscribe","symbols":["AAPL","GOOGL"]}
Server sends: {"type":"quote","symbol":"AAPL","bid":150.23,"ask":150.25,
               "last":150.24,"volume":12345,"timestamp":<unix_us>}
"""

import asyncio
import json
import time
import random
import sys
import struct
import hashlib
import base64

# ── WebSocket implementation (RFC 6455, no dependencies) ─────────

class WebSocketClient:
    def __init__(self, reader, writer):
        self.reader = reader
        self.writer = writer
        self.subscriptions = set()

    async def send_frame(self, payload: bytes, opcode: int = 0x1):
        """Send unmasked text frame (server→client)."""
        frame = bytearray()
        frame.append(0x80 | opcode)  # FIN=1
        length = len(payload)
        if length < 126:
            frame.append(length)
        elif length < 65536:
            frame.append(126)
            frame.extend(struct.pack('>H', length))
        else:
            frame.append(127)
            frame.extend(struct.pack('>Q', length))
        frame.extend(payload)
        self.writer.write(bytes(frame))
        await self.writer.drain()

    async def recv_frame(self):
        """Receive masked frame (client→server). Returns (opcode, payload)."""
        hdr = await self.reader.readexactly(2)
        opcode = hdr[0] & 0x0F
        masked = (hdr[1] >> 7) & 1
        plen = hdr[1] & 0x7F

        if plen == 126:
            ext = await self.reader.readexactly(2)
            plen = struct.unpack('>H', ext)[0]
        elif plen == 127:
            ext = await self.reader.readexactly(8)
            plen = struct.unpack('>Q', ext)[0]

        mask_key = b''
        if masked:
            mask_key = await self.reader.readexactly(4)

        payload = await self.reader.readexactly(plen)
        if masked:
            payload = bytes(b ^ mask_key[i % 4] for i, b in enumerate(payload))

        return opcode, payload

    async def close(self):
        try:
            self.writer.close()
            await self.writer.wait_closed()
        except Exception:
            pass


# ── Market data simulator ────────────────────────────────────────

SYMBOLS = ['AAPL', 'GOOGL', 'MSFT', 'TSLA', 'AMZN', 'NVDA', 'META', 'JPM']
BASE_PRICES = {
    'AAPL': 185.0, 'GOOGL': 140.0, 'MSFT': 420.0, 'TSLA': 245.0,
    'AMZN': 185.0, 'NVDA': 880.0, 'META': 510.0, 'JPM': 200.0,
}

class MarketDataServer:
    def __init__(self, push_interval_ms: int = 100):
        self.push_interval_ms = push_interval_ms
        self.clients: list[WebSocketClient] = []
        self.symbol_subscribers: dict[str, set[WebSocketClient]] = {
            s: set() for s in SYMBOLS
        }
        self.prices = dict(BASE_PRICES)
        self.volumes = {s: random.randint(1000, 100000) for s in SYMBOLS}

    def add_client(self, client: WebSocketClient):
        self.clients.append(client)

    def remove_client(self, client: WebSocketClient):
        self.clients = [c for c in self.clients if c is not client]
        for sym_set in self.symbol_subscribers.values():
            sym_set.discard(client)

    def subscribe(self, client: WebSocketClient, symbols: list[str]):
        for sym in symbols:
            if sym in self.symbol_subscribers:
                self.symbol_subscribers[sym].add(client)
                client.subscriptions.add(sym)

    def unsubscribe(self, client: WebSocketClient, symbols: list[str]):
        for sym in symbols:
            if sym in self.symbol_subscribers:
                self.symbol_subscribers[sym].discard(client)
                client.subscriptions.discard(sym)

    def tick_price(self, symbol: str):
        """Random walk price movement."""
        volatility = BASE_PRICES[symbol] * 0.002  # 0.2% volatility
        change = random.gauss(0, volatility)
        self.prices[symbol] = max(0.01, self.prices[symbol] + change)
        spread = self.prices[symbol] * 0.0001  # 0.01% spread
        bid = round(self.prices[symbol] - spread / 2, 2)
        ask = round(self.prices[symbol] + spread / 2, 2)
        last = round(self.prices[symbol], 2)
        self.volumes[symbol] += random.randint(0, 500)
        return bid, ask, last

    async def push_loop(self):
        """Continuously push market data to subscribed clients."""
        seq = 0
        while True:
            now_us = int(time.time() * 1_000_000)
            for symbol in SYMBOLS:
                subs = self.symbol_subscribers.get(symbol, set())
                if not subs:
                    continue

                bid, ask, last = self.tick_price(symbol)
                msg = json.dumps({
                    'type': 'quote',
                    'symbol': symbol,
                    'bid': bid,
                    'ask': ask,
                    'last': last,
                    'volume': self.volumes[symbol],
                    'timestamp': now_us,
                    'seq': seq,
                })
                seq += 1

                payload = msg.encode()
                dead = []
                for client in list(subs):
                    try:
                        await client.send_frame(payload)
                    except Exception:
                        dead.append(client)

                for d in dead:
                    self.remove_client(d)

            await asyncio.sleep(self.push_interval_ms / 1000.0)


# ── WebSocket handshake ──────────────────────────────────────────

def compute_accept(key: str) -> str:
    GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
    sha1 = hashlib.sha1((key + GUID).encode()).digest()
    return base64.b64encode(sha1).decode()


async def handle_client(reader, writer, server: MarketDataServer):
    """Handle one WebSocket client connection."""
    ws = WebSocketClient(reader, writer)

    try:
        # Read HTTP upgrade request
        request = b''
        while b'\r\n\r\n' not in request:
            chunk = await reader.read(4096)
            if not chunk:
                return
            request += chunk

        req_text = request.decode(errors='replace')
        ws_key = ''
        for line in req_text.split('\r\n'):
            if line.lower().startswith('sec-websocket-key:'):
                ws_key = line.split(':', 1)[1].strip()
                break

        if not ws_key:
            return

        # Send 101 response
        accept = compute_accept(ws_key)
        response = (
            'HTTP/1.1 101 Switching Protocols\r\n'
            'Upgrade: websocket\r\n'
            'Connection: Upgrade\r\n'
            f'Sec-WebSocket-Accept: {accept}\r\n'
            '\r\n'
        )
        writer.write(response.encode())
        await writer.drain()

        server.add_client(ws)
        addr = writer.get_extra_info('peername')
        print(f'[CONNECT] {addr}')

        # Message loop
        while True:
            opcode, payload = await ws.recv_frame()

            if opcode == 0x8:  # Close
                print(f'[CLOSE] {addr}')
                await ws.send_frame(b'', 0x8)
                break
            elif opcode == 0x9:  # Ping
                await ws.send_frame(payload, 0xA)  # Pong
                continue
            elif opcode == 0xA:  # Pong
                continue

            # Text message — parse JSON
            try:
                data = json.loads(payload.decode())
            except (json.JSONDecodeError, UnicodeDecodeError):
                continue

            action = data.get('action', '')
            if action == 'subscribe':
                symbols = data.get('symbols', [])
                server.subscribe(ws, symbols)
                resp = json.dumps({
                    'status': 'ok',
                    'action': 'subscribed',
                    'symbols': symbols,
                    'timestamp': int(time.time() * 1_000_000),
                })
                await ws.send_frame(resp.encode())
                print(f'[SUB] {addr} → {symbols}')
            elif action == 'unsubscribe':
                symbols = data.get('symbols', [])
                server.unsubscribe(ws, symbols)
                print(f'[UNSUB] {addr} → {symbols}')

    except (ConnectionError, asyncio.IncompleteReadError, OSError):
        pass
    finally:
        server.remove_client(ws)
        await ws.close()


# ── Main ─────────────────────────────────────────────────────────

async def main(port: int = 8080, push_interval_ms: int = 100):
    server = MarketDataServer(push_interval_ms)

    # Start price push loop
    push_task = asyncio.create_task(server.push_loop())

    # Start TCP server
    async def on_connect(reader, writer):
        await handle_client(reader, writer, server)

    tcp_server = await asyncio.start_server(on_connect, '0.0.0.0', port)

    print(f'╔══════════════════════════════════════╗')
    print(f'║  Market Data Server                 ║')
    print(f'╠══════════════════════════════════════╣')
    print(f'║  Port: {port}                         ║')
    print(f'║  Symbols: {", ".join(SYMBOLS[:4])}...      ║')
    print(f'║  Push: every {push_interval_ms}ms               ║')
    print(f'╚══════════════════════════════════════╝')
    print()
    print(f'Subscribe: {{"action":"subscribe","symbols":["AAPL","GOOGL"]}}')

    try:
        await tcp_server.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        push_task.cancel()
        tcp_server.close()
        await tcp_server.wait_closed()


if __name__ == '__main__':
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 8080
    interval = int(sys.argv[2]) if len(sys.argv) > 2 else 100
    asyncio.run(main(port, interval))