kite_client.py

import json
import os
import threading
import time as time_mod
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Tuple

import pandas as pd
import requests
from kiteconnect import KiteConnect, KiteTicker
from kiteconnect.exceptions import TokenException

from config import KITE_API_KEY

BASE_DIR = Path(__file__).resolve().parent
TOKEN_PATH = BASE_DIR / "tokens.json"
NSE_PATH = BASE_DIR / "instruments_nse.csv"
NFO_PATH = BASE_DIR / "instruments_nfo.csv"

_KITE_CACHE = None
_KITE_CACHE_TOKEN = None
_TICKER_MANAGER = None
_INSTRUMENT_TOKEN_MAP = None

RELAY_BASE_URL = os.getenv("ORDER_RELAY_BASE_URL", "").strip().rstrip("/")
RELAY_SHARED_SECRET = os.getenv("ORDER_RELAY_SHARED_SECRET", "").strip()
RELAY_TIMEOUT_SECONDS = float(os.getenv("ORDER_RELAY_TIMEOUT_SECONDS", "15"))


def _resolve_access_token():
    access_token = os.getenv("KITE_ACCESS_TOKEN")

    if not access_token:
        if not TOKEN_PATH.exists():
            raise FileNotFoundError(
                "No KITE_ACCESS_TOKEN env var and tokens.json not found. "
                "Run main.py locally or set KITE_ACCESS_TOKEN."
            )
        with open(TOKEN_PATH, "r") as f:
            token_data = json.load(f)
        access_token = token_data.get("access_token")

    if not access_token:
        raise ValueError("access_token missing. Set KITE_ACCESS_TOKEN or regenerate tokens.json")

    return access_token


def get_kite(force_refresh: bool = False):
    global _KITE_CACHE, _KITE_CACHE_TOKEN

    access_token = _resolve_access_token()
    if not force_refresh and _KITE_CACHE is not None and _KITE_CACHE_TOKEN == access_token:
        return _KITE_CACHE

    kite = KiteConnect(api_key=KITE_API_KEY)
    kite.set_access_token(access_token)

    try:
        kite.profile()
    except TokenException as e:
        raise RuntimeError(
            "Kite access token is invalid or expired. "
            "Generate a fresh token and update KITE_ACCESS_TOKEN."
        ) from e

    _KITE_CACHE = kite
    _KITE_CACHE_TOKEN = access_token
    return _KITE_CACHE


def _load_instrument_token_map() -> Dict[str, int]:
    global _INSTRUMENT_TOKEN_MAP
    if _INSTRUMENT_TOKEN_MAP is not None:
        return dict(_INSTRUMENT_TOKEN_MAP)

    mapping: Dict[str, int] = {}
    for exchange, path in (("NSE", NSE_PATH), ("NFO", NFO_PATH)):
        if not path.exists():
            continue
        df = pd.read_csv(path, usecols=["tradingsymbol", "instrument_token"])
        df["tradingsymbol"] = df["tradingsymbol"].astype(str).str.upper()
        df["instrument_token"] = pd.to_numeric(df["instrument_token"], errors="coerce")
        df = df.dropna(subset=["instrument_token"])
        for _, row in df.iterrows():
            mapping[f"{exchange}:{row['tradingsymbol']}"] = int(row["instrument_token"])

    _INSTRUMENT_TOKEN_MAP = mapping
    return dict(_INSTRUMENT_TOKEN_MAP)


class LiveTickerManager:
    def __init__(self):
        self._lock = threading.RLock()
        self._started = False
        self._connected = False
        self._last_error = ""
        self._kws: Optional[KiteTicker] = None
        self._latest_ticks: Dict[int, dict] = {}
        self._token_to_key: Dict[int, str] = {}
        self._key_to_token: Dict[str, int] = {}
        self._token_mode: Dict[int, str] = {}
        self._instrument_tokens = _load_instrument_token_map()
        self._connected_event = threading.Event()

    def start(self):
        with self._lock:
            if self._started and self._kws is not None:
                return self

            access_token = _resolve_access_token()
            self._kws = KiteTicker(
                KITE_API_KEY,
                access_token,
                reconnect=True,
                reconnect_max_tries=300,
                reconnect_max_delay=60,
            )
            self._kws.on_ticks = self._on_ticks
            self._kws.on_connect = self._on_connect
            self._kws.on_close = self._on_close
            self._kws.on_error = self._on_error
            self._kws.on_reconnect = self._on_reconnect
            self._kws.on_noreconnect = self._on_noreconnect
            self._kws.connect(threaded=True)
            self._started = True
            return self

    def _mode_constant(self, mode: str):
        mode = str(mode or "ltp").lower()
        if mode == "full":
            return self._kws.MODE_FULL
        if mode == "quote":
            return self._kws.MODE_QUOTE
        return self._kws.MODE_LTP

    def _apply_subscriptions(self, token_mode: Dict[int, str]):
        if not token_mode or self._kws is None:
            return
        tokens = sorted(token_mode.keys())
        self._kws.subscribe(tokens)
        by_mode: Dict[str, List[int]] = {}
        for token, mode in token_mode.items():
            by_mode.setdefault(mode, []).append(token)
        for mode, mode_tokens in by_mode.items():
            self._kws.set_mode(self._mode_constant(mode), sorted(mode_tokens))

    def _on_connect(self, ws, response):
        with self._lock:
            self._connected = True
            self._last_error = ""
            token_mode = dict(self._token_mode)
        self._connected_event.set()
        self._apply_subscriptions(token_mode)

    def _on_ticks(self, ws, ticks):
        with self._lock:
            for tick in ticks or []:
                token = int(tick.get("instrument_token"))
                tick = dict(tick)
                tick["received_at"] = pd.Timestamp.now(tz="Asia/Kolkata")
                self._latest_ticks[token] = tick

    def _on_close(self, ws, code, reason):
        with self._lock:
            self._connected = False
            self._last_error = f"close:{code}:{reason}"
        self._connected_event.clear()

    def _on_error(self, ws, code, reason):
        with self._lock:
            self._connected = False
            self._last_error = f"error:{code}:{reason}"
        self._connected_event.clear()

    def _on_reconnect(self, ws, attempts_count):
        with self._lock:
            self._last_error = f"reconnecting:{attempts_count}"

    def _on_noreconnect(self, ws):
        with self._lock:
            self._connected = False
            self._last_error = "reconnect_exhausted"
        self._connected_event.clear()

    def subscribe_symbols(self, symbols: Iterable[Tuple[str, str]], mode: str = "ltp") -> int:
        self.start()
        mode = str(mode or "ltp").lower()
        added: Dict[int, str] = {}

        with self._lock:
            for exchange, tradingsymbol in symbols or []:
                key = f"{str(exchange).upper()}:{str(tradingsymbol).upper()}"
                token = self._instrument_tokens.get(key)
                if token is None:
                    continue
                self._token_to_key[token] = key
                self._key_to_token[key] = token
                existing_mode = self._token_mode.get(token)
                desired_mode = mode
                if existing_mode == "full" or desired_mode == existing_mode:
                    continue
                if existing_mode == "quote" and desired_mode == "ltp":
                    continue
                self._token_mode[token] = desired_mode
                added[token] = desired_mode
                if existing_mode is None:
                    continue
            for exchange, tradingsymbol in symbols or []:
                key = f"{str(exchange).upper()}:{str(tradingsymbol).upper()}"
                token = self._instrument_tokens.get(key)
                if token is None:
                    continue
                if token not in self._token_mode:
                    self._token_mode[token] = mode
                    added[token] = mode
                    self._token_to_key[token] = key
                    self._key_to_token[key] = token

            connected = self._connected
            kws = self._kws

        if connected and kws is not None and added:
            self._apply_subscriptions(added)

        return len(added)

    def get_tick(self, exchange: str, tradingsymbol: str) -> Optional[dict]:
        key = f"{str(exchange).upper()}:{str(tradingsymbol).upper()}"
        with self._lock:
            token = self._key_to_token.get(key)
            if token is None:
                return None
            tick = self._latest_ticks.get(token)
            return dict(tick) if tick is not None else None

    def get_last_price(self, exchange: str, tradingsymbol: str) -> Optional[float]:
        tick = self.get_tick(exchange, tradingsymbol)
        if not tick:
            return None
        try:
            return float(tick.get("last_price"))
        except Exception:
            return None

    def status(self) -> dict:
        with self._lock:
            return {
                "started": self._started,
                "connected": self._connected,
                "subscribed_tokens": len(self._token_mode),
                "ticks_cached": len(self._latest_ticks),
                "last_error": self._last_error,
            }


def get_live_ticker_manager(force_refresh: bool = False) -> LiveTickerManager:
    global _TICKER_MANAGER
    if force_refresh or _TICKER_MANAGER is None:
        _TICKER_MANAGER = LiveTickerManager().start()
    return _TICKER_MANAGER


def get_quote_map(instruments):
    instruments = [str(x).strip().upper() for x in (instruments or []) if str(x).strip()]
    if not instruments:
        return {}
    kite = get_kite()
    return kite.quote(instruments)


def get_ltp_map(instruments):
    instruments = [str(x).strip().upper() for x in (instruments or []) if str(x).strip()]
    if not instruments:
        return {}
    kite = get_kite()
    return kite.ltp(instruments)



def _normalize_order_snapshot(order: Optional[dict]) -> dict:
    order = dict(order or {})
    return {
        "order_id": order.get("order_id"),
        "status": str(order.get("status") or "").upper(),
        "status_message": order.get("status_message") or order.get("status_message_raw") or "",
        "average_price": order.get("average_price"),
        "filled_quantity": order.get("filled_quantity"),
        "pending_quantity": order.get("pending_quantity"),
        "exchange_update_timestamp": order.get("exchange_update_timestamp") or order.get("order_timestamp"),
        "tradingsymbol": order.get("tradingsymbol"),
        "transaction_type": order.get("transaction_type"),
        "quantity": order.get("quantity"),
        "price": order.get("price"),
        "raw": order,
    }


def _relay_enabled() -> bool:
    return bool(RELAY_BASE_URL)


def _relay_headers() -> dict:
    if not RELAY_SHARED_SECRET:
        raise RuntimeError(
            "ORDER_RELAY_SHARED_SECRET missing. Set it in HF backend env before using relay order routing."
        )
    return {"X-Relay-Secret": RELAY_SHARED_SECRET}


def _relay_post(path: str, payload: dict, timeout_seconds: Optional[float] = None) -> dict:
    if not _relay_enabled():
        raise RuntimeError(
            "ORDER_RELAY_BASE_URL missing. Set it to your AWS relay URL before using relay order routing."
        )
    timeout = float(timeout_seconds or RELAY_TIMEOUT_SECONDS)
    url = f"{RELAY_BASE_URL}{path}"
    try:
        resp = requests.post(url, json=payload, headers=_relay_headers(), timeout=timeout)
    except Exception as e:
        raise RuntimeError(f"order relay request failed: {e}") from e

    try:
        body = resp.json()
    except Exception:
        body = {"detail": resp.text}

    if resp.status_code >= 400:
        raise RuntimeError(f"order relay error {resp.status_code}: {body}")
    return dict(body or {})


def get_order_snapshot(order_id: str) -> dict:
    if _relay_enabled():
        return _relay_post("/order-snapshot", {"order_id": order_id})

    kite = get_kite()
    try:
        history = kite.order_history(order_id)
        if history:
            return _normalize_order_snapshot(history[-1])
    except Exception:
        pass
    try:
        orders = kite.orders()
        for order in reversed(orders or []):
            if str(order.get("order_id")) == str(order_id):
                return _normalize_order_snapshot(order)
    except Exception:
        pass
    return _normalize_order_snapshot({"order_id": order_id, "status": "UNKNOWN"})


def wait_for_order_terminal(order_id: str, timeout_seconds: float = 8.0, poll_interval_seconds: float = 0.75) -> dict:
    if _relay_enabled():
        return _relay_post(
            "/wait-terminal",
            {
                "order_id": order_id,
                "timeout_seconds": float(timeout_seconds),
                "poll_interval_seconds": float(poll_interval_seconds),
            },
            timeout_seconds=max(float(timeout_seconds), RELAY_TIMEOUT_SECONDS) + 5.0,
        )

    deadline = time_mod.time() + max(float(timeout_seconds), 0.0)
    latest = get_order_snapshot(order_id)
    terminal = {"COMPLETE", "CANCELLED", "REJECTED"}
    while time_mod.time() <= deadline:
        latest = get_order_snapshot(order_id)
        if latest.get("status") in terminal:
            return latest
        time_mod.sleep(max(float(poll_interval_seconds), 0.2))
    return latest


def cancel_order(order_id: str, variety: str = "regular") -> dict:
    if _relay_enabled():
        return _relay_post("/cancel-order", {"order_id": order_id, "variety": variety})

    kite = get_kite()
    try:
        kite.cancel_order(variety=variety, order_id=order_id)
    except Exception:
        pass
    return get_order_snapshot(order_id)


def place_fno_limit_order(
    tradingsymbol: str,
    transaction_type: str,
    quantity: int,
    limit_price: float,
    product: str = "MIS",
    variety: str = "regular",
    exchange: str = "NFO",
    tag: Optional[str] = None,
) -> dict:
    tx = str(transaction_type or "").upper()
    if tx not in {"BUY", "SELL"}:
        raise ValueError(f"Unsupported transaction_type: {transaction_type}")

    if _relay_enabled():
        return _relay_post(
            "/place-order",
            {
                "tradingsymbol": str(tradingsymbol).upper(),
                "transaction_type": tx,
                "quantity": int(quantity),
                "limit_price": float(limit_price),
                "product": product,
                "variety": variety,
                "exchange": exchange,
                "tag": (str(tag)[:20] if tag else None),
            },
        )

    kite = get_kite()
    order_id = kite.place_order(
        variety=variety,
        exchange=exchange,
        tradingsymbol=str(tradingsymbol).upper(),
        transaction_type=tx,
        quantity=int(quantity),
        product=product,
        order_type=kite.ORDER_TYPE_LIMIT,
        price=float(limit_price),
        validity=kite.VALIDITY_DAY,
        tag=(str(tag)[:20] if tag else None),
    )
    snapshot = get_order_snapshot(order_id)
    snapshot["order_id"] = order_id
    return snapshot


def get_gtt(trigger_id) -> dict:
    """Fetch one GTT (status + legs). status is one of active/triggered/
    cancelled/expired/disabled/deleted."""
    if _relay_enabled():
        return _relay_post("/get-gtt", {"trigger_id": int(trigger_id)})
    return get_kite().get_gtt(int(trigger_id)) or {}


def cancel_gtt(trigger_id) -> dict:
    if _relay_enabled():
        return _relay_post("/cancel-gtt", {"trigger_id": int(trigger_id)})
    return get_kite().delete_gtt(int(trigger_id))


def get_net_position_qty(tradingsymbol: str, exchange: str = "NFO") -> int:
    """Net held quantity for a symbol (0 = flat). Used to confirm we still hold
    a position before the bot places a square-off sell (a GTT may have already
    closed it — selling again would open a SHORT)."""
    sym = str(tradingsymbol).upper()
    if _relay_enabled():
        data = _relay_post("/positions", {})
        nets = data.get("net", data) if isinstance(data, dict) else data
    else:
        nets = get_kite().positions().get("net", [])
    total = 0
    for p in nets or []:
        if str(p.get("tradingsymbol", "")).upper() == sym and \
                str(p.get("exchange", exchange)).upper() == exchange.upper():
            total += int(p.get("quantity", 0) or 0)
    return total


def find_completed_sell(tradingsymbol: str) -> Optional[dict]:
    """Most recent COMPLETE SELL of ``tradingsymbol`` in today's orderbook —
    used to recover the actual fill of a GTT-triggered exit."""
    sym = str(tradingsymbol).upper()
    if _relay_enabled():
        data = _relay_post("/orders", {})
        orders = data.get("orders", data) if isinstance(data, dict) else data
    else:
        orders = get_kite().orders()
    hits = [o for o in (orders or [])
            if str(o.get("tradingsymbol", "")).upper() == sym
            and str(o.get("transaction_type", "")).upper() == "SELL"
            and str(o.get("status", "")).upper() == "COMPLETE"]
    if not hits:
        return None
    o = hits[-1]
    return {"average_price": _safe_float(o.get("average_price")),
            "order_id": o.get("order_id"), "timestamp": o.get("order_timestamp")}


def _safe_float(v):
    try:
        return float(v)
    except (TypeError, ValueError):
        return 0.0


def place_gtt_oco_sell(
    tradingsymbol: str,
    quantity: int,
    last_price: float,
    target_price: float,
    stop_price: float,
    exchange: str = "NFO",
    product: str = "MIS",
    stop_limit_buffer_pct: float = 0.0,
) -> dict:
    """Place a two-leg (OCO) GTT to bracket a long option position: an upper
    leg that SELLs at ``target_price`` and a lower leg that SELLs at
    ``stop_price``. When either trigger hits, Kite places that leg's order and
    cancels the other. Broker-held, so it exits without bot latency.

    ``last_price`` must be the option's current LTP (Kite validates triggers
    against it: lower trigger < last < upper trigger). Returns {"trigger_id": ...}.

    NOTE: Zerodha GTT places only LIMIT orders (SL-M / MARKET are NOT supported
    in a GTT). The target leg is a LIMIT at ``target_price`` (good). The stop leg
    triggers at ``stop_price`` but, to stay fillable in a fast drop, its LIMIT
    price is set ``stop_limit_buffer_pct`` % BELOW the trigger (0 = limit exactly
    at the trigger, which can be skipped past in a gap; ~2-3 makes it marketable
    at a small extra-slippage cost).
    """
    sym = str(tradingsymbol).upper()
    qty = int(quantity)
    target_price = round(float(target_price), 1)
    stop_price = round(float(stop_price), 1)
    if not (stop_price < float(last_price) < target_price):
        raise ValueError(
            f"GTT OCO needs stop({stop_price}) < last({last_price}) < target({target_price})")

    stop_limit = round(stop_price * (1 - max(0.0, stop_limit_buffer_pct) / 100.0), 1)
    orders = [
        {"exchange": exchange, "tradingsymbol": sym, "transaction_type": "SELL",
         "quantity": qty, "order_type": "LIMIT", "product": product, "price": stop_limit},
        {"exchange": exchange, "tradingsymbol": sym, "transaction_type": "SELL",
         "quantity": qty, "order_type": "LIMIT", "product": product, "price": target_price},
    ]
    payload = {
        "trigger_type": "two-leg", "tradingsymbol": sym, "exchange": exchange,
        "trigger_values": [stop_price, target_price], "last_price": float(last_price),
        "orders": orders,
    }
    if _relay_enabled():
        return _relay_post("/place-gtt", payload)

    kite = get_kite()
    res = kite.place_gtt(
        trigger_type=kite.GTT_TYPE_OCO, tradingsymbol=sym, exchange=exchange,
        trigger_values=[stop_price, target_price], last_price=float(last_price),
        orders=orders,
    )
    return {"trigger_id": res.get("trigger_id") if isinstance(res, dict) else res}