data_feed.py

"""
data_feed.py
------------
Fetches OHLCV candle data from Hyperliquid via REST API.
Maintains a live candle buffer updated via WebSocket ticks.

Responsibilities:
- REST: fetch historical candles for indicator seeding.
- WebSocket: subscribe to live trade ticks and aggregate into candles.
- Expose a clean get_candles(symbol, timeframe, limit) interface.
"""

import time
import threading
import logging
from datetime import datetime, timezone
from typing import Optional

import pandas as pd
import requests
import websocket
import json

from config import Config

logger = logging.getLogger(__name__)


class CandleBuffer:
    """Thread-safe rolling buffer of OHLCV candles for a single timeframe."""

    def __init__(self, timeframe_minutes: int, max_candles: int = 200):
        self.tf_minutes = timeframe_minutes
        self.tf_ms = timeframe_minutes * 60 * 1000
        self.max_candles = max_candles
        self._candles: list[dict] = []
        self._lock = threading.Lock()

    def update_from_rest(self, candles: list[dict]) -> None:
        """Seed the buffer with REST candle data."""
        with self._lock:
            self._candles = candles[-self.max_candles:]
        logger.debug(f"[CandleBuffer TF={self.tf_minutes}m] Seeded {len(self._candles)} candles.")

    def update_tick(self, price: float, volume: float, timestamp_ms: int) -> None:
        """Incorporate a live trade tick into the current open candle."""
        candle_open_ts = (timestamp_ms // self.tf_ms) * self.tf_ms

        with self._lock:
            if self._candles and self._candles[-1]["t"] == candle_open_ts:
                # Update existing open candle
                c = self._candles[-1]
                c["h"] = max(c["h"], price)
                c["l"] = min(c["l"], price)
                c["c"] = price
                c["v"] += volume
            else:
                # Open a new candle
                new_candle = {
                    "t": candle_open_ts,
                    "o": price,
                    "h": price,
                    "l": price,
                    "c": price,
                    "v": volume,
                }
                self._candles.append(new_candle)
                if len(self._candles) > self.max_candles:
                    self._candles.pop(0)

    def get_dataframe(self) -> pd.DataFrame:
        """Return candles as a pandas DataFrame."""
        with self._lock:
            if not self._candles:
                return pd.DataFrame()
            df = pd.DataFrame(self._candles)
        df.rename(columns={"t": "timestamp", "o": "open", "h": "high",
                            "l": "low", "c": "close", "v": "volume"}, inplace=True)
        df["timestamp"] = pd.to_datetime(df["timestamp"], unit="ms", utc=True)
        df.set_index("timestamp", inplace=True)
        return df


class DataFeed:
    """
    Main data feed class.
    - Fetches historical candles via Hyperliquid REST.
    - Subscribes to live WebSocket ticks to keep candles current.
    """

    def __init__(self, config: Config):
        self.config = config
        self.symbol = config.SYMBOL
        self.rest_url = config.HL_REST_URL
        self.ws_url = config.HL_WS_URL

        # One buffer per timeframe
        self.buffers: dict[int, CandleBuffer] = {
            config.ENTRY_TF: CandleBuffer(config.ENTRY_TF),
            config.PRIMARY_TF: CandleBuffer(config.PRIMARY_TF),
            config.FILTER_TF: CandleBuffer(config.FILTER_TF),
        }

        self._ws: Optional[websocket.WebSocketApp] = None
        self._ws_thread: Optional[threading.Thread] = None
        self._running = False

    # ------------------------------------------------------------------
    # REST Methods
    # ------------------------------------------------------------------

    def fetch_candles_rest(self, timeframe_minutes: int, limit: int = 100) -> list[dict]:
        """
        Fetch historical OHLCV candles from Hyperliquid REST API.
        Returns a list of candle dicts with keys: t, o, h, l, c, v
        """
        interval_map = {1: "1m", 3: "3m", 5: "5m", 15: "15m",
                        30: "30m", 60: "1h", 240: "4h", 1440: "1d"}
        interval = interval_map.get(timeframe_minutes)
        if not interval:
            raise ValueError(f"Unsupported timeframe: {timeframe_minutes}m")

        end_time = int(time.time() * 1000)
        start_time = end_time - (timeframe_minutes * 60 * 1000 * limit)

        payload = {
            "type": "candleSnapshot",
            "req": {
                "coin": self.symbol,
                "interval": interval,
                "startTime": start_time,
                "endTime": end_time,
            }
        }

        try:
            resp = requests.post(
                f"{self.rest_url}/info",
                json=payload,
                timeout=10
            )
            resp.raise_for_status()
            raw = resp.json()

            candles = []
            for c in raw:
                candles.append({
                    "t": int(c["t"]),
                    "o": float(c["o"]),
                    "h": float(c["h"]),
                    "l": float(c["l"]),
                    "c": float(c["c"]),
                    "v": float(c["v"]),
                })
            logger.info(f"[DataFeed] Fetched {len(candles)} candles for {self.symbol} @ {timeframe_minutes}m")
            return candles

        except requests.RequestException as e:
            logger.error(f"[DataFeed] REST fetch failed for {timeframe_minutes}m: {e}")
            return []

    def seed_all_buffers(self) -> None:
        """Fetch historical candles for all timeframes and seed buffers."""
        for tf, buffer in self.buffers.items():
            candles = self.fetch_candles_rest(tf, limit=self.config.CANDLE_LOOKBACK)
            if candles:
                buffer.update_from_rest(candles)
            else:
                logger.warning(f"[DataFeed] No candles returned for TF={tf}m — buffer empty.")

    # ------------------------------------------------------------------
    # WebSocket Methods
    # ------------------------------------------------------------------

    def _on_ws_message(self, ws, message: str) -> None:
        """Handle incoming WebSocket messages and route ticks to buffers."""
        try:
            data = json.loads(message)

            # Hyperliquid WS trade message format
            if data.get("channel") == "trades":
                for trade in data.get("data", []):
                    price = float(trade["px"])
                    volume = float(trade["sz"])
                    ts_ms = int(trade["time"])

                    for buffer in self.buffers.values():
                        buffer.update_tick(price, volume, ts_ms)

        except (json.JSONDecodeError, KeyError, ValueError) as e:
            logger.debug(f"[DataFeed] WS message parse error: {e}")

    def _on_ws_open(self, ws) -> None:
        """Subscribe to trade stream on WebSocket open."""
        sub_msg = {
            "method": "subscribe",
            "subscription": {
                "type": "trades",
                "coin": self.symbol
            }
        }
        ws.send(json.dumps(sub_msg))
        logger.info(f"[DataFeed] WebSocket subscribed to trades for {self.symbol}")

    def _on_ws_error(self, ws, error) -> None:
        logger.error(f"[DataFeed] WebSocket error: {error}")

    def _on_ws_close(self, ws, close_status_code, close_msg) -> None:
        logger.warning(f"[DataFeed] WebSocket closed: {close_status_code} — {close_msg}")
        if self._running:
            logger.info("[DataFeed] Attempting WebSocket reconnect in 5s...")
            time.sleep(5)
            self._start_websocket()

    def _start_websocket(self) -> None:
        """Initialize and start the WebSocket connection in a daemon thread."""
        self._ws = websocket.WebSocketApp(
            self.ws_url,
            on_open=self._on_ws_open,
            on_message=self._on_ws_message,
            on_error=self._on_ws_error,
            on_close=self._on_ws_close,
        )
        self._ws_thread = threading.Thread(
            target=self._ws.run_forever,
            kwargs={"ping_interval": 30, "ping_timeout": 10},
            daemon=True,
            name="ws-data-feed"
        )
        self._ws_thread.start()

    # ------------------------------------------------------------------
    # Public Interface
    # ------------------------------------------------------------------

    def start(self) -> None:
        """Seed buffers from REST then start live WebSocket feed."""
        self._running = True
        logger.info("[DataFeed] Seeding candle buffers from REST...")
        self.seed_all_buffers()
        logger.info("[DataFeed] Starting WebSocket live feed...")
        self._start_websocket()

    def stop(self) -> None:
        """Gracefully stop the WebSocket feed."""
        self._running = False
        if self._ws:
            self._ws.close()
        logger.info("[DataFeed] DataFeed stopped.")

    def get_candles(self, timeframe_minutes: int) -> pd.DataFrame:
        """
        Return the current candle DataFrame for a given timeframe.
        Returns empty DataFrame if buffer is not seeded yet.
        """
        buffer = self.buffers.get(timeframe_minutes)
        if buffer is None:
            raise ValueError(f"No buffer registered for TF={timeframe_minutes}m")
        return buffer.get_dataframe()

    def get_latest_price(self) -> Optional[float]:
        """Return the latest close price from the entry timeframe buffer."""
        df = self.get_candles(self.config.ENTRY_TF)
        if df.empty:
            return None
        return float(df["close"].iloc[-1])