"""
Hyperliquid BTC perpetual trading client.

KEY CONCEPT — two wallets:
  • account_address  : Your MetaMask address. Holds USDC and open positions.
                       Used for all info/query calls.
  • api_private_key  : API wallet private key generated at app.hyperliquid.xyz/API.
                       Used only for signing orders. Cannot withdraw.

Both are read from config (hl_account_address, hl_api_private_key).
"""

import asyncio
import logging
import time
from functools import partial
from typing import Optional

from eth_account import Account
from hyperliquid.exchange import Exchange
from hyperliquid.info import Info

logger = logging.getLogger(__name__)

HL_MAINNET_URL  = "https://api.hyperliquid.xyz"
HL_TESTNET_URL  = "https://api.hyperliquid-testnet.xyz"

# Precision map — Hyperliquid's szDecimals for common coins
# Fetched dynamically via meta() but cached here as fallback
SZ_DECIMALS_FALLBACK = {"BTC": 5, "ETH": 4}

# Module-level cache: coin → (szDecimals, expiry_timestamp)
_SZ_DECIMALS_CACHE: dict[str, tuple[int, float]] = {}
_SZ_DECIMALS_TTL = 300  # 5 minutes

# Max-leverage cache: coin → (maxLeverage, expiry_timestamp).
# meta() is a full-universe call (~1 KB JSON), cheap but wasteful on every
# trade open when the leverage caps change at most once a month. 5-min TTL
# matches _SZ_DECIMALS_TTL so a single meta() response can refresh both.
_MAX_LEV_CACHE: dict[str, tuple[int, float]] = {}
_MAX_LEV_TTL = 300  # 5 minutes


class HyperliquidTrader:
    def __init__(
        self,
        api_private_key: str,
        account_address: str,
        leverage: int = 3,
        mainnet: bool = True,
    ):
        """
        Args:
            api_private_key:  Private key of the API wallet (from Hyperliquid dashboard).
            account_address:  Main MetaMask wallet address (holds USDC + positions).
            leverage:         Isolated leverage to set before each trade.
            mainnet:          True = mainnet, False = testnet.
        """
        self._api_wallet = Account.from_key(api_private_key)
        self._account_address = account_address.lower()
        self._leverage = leverage
        self._base_url = HL_MAINNET_URL if mainnet else HL_TESTNET_URL

        # Exchange signs with api_wallet but acts on behalf of account_address
        self._exchange = Exchange(
            self._api_wallet,
            self._base_url,
            account_address=self._account_address,
        )
        self._info = Info(self._base_url, skip_ws=True)

    # ── internal helpers ──────────────────────────────────────────────────────

    async def _run(self, fn, *args, **kwargs):
        """Run blocking SDK calls in a thread pool without blocking the event loop."""
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(None, partial(fn, *args, **kwargs))

    async def _get_sz_decimals(self, coin: str) -> int:
        now = time.monotonic()
        cached = _SZ_DECIMALS_CACHE.get(coin)
        if cached is not None and now < cached[1]:
            return cached[0]
        try:
            meta = await self._run(self._info.meta)
            for u in meta.get("universe", []):
                name = u.get("name")
                val = int(u.get("szDecimals", SZ_DECIMALS_FALLBACK.get(name, 4)))
                _SZ_DECIMALS_CACHE[name] = (val, now + _SZ_DECIMALS_TTL)
            if coin in _SZ_DECIMALS_CACHE:
                return _SZ_DECIMALS_CACHE[coin][0]
        except Exception:
            pass
        return SZ_DECIMALS_FALLBACK.get(coin, 4)

    async def _get_max_leverage(self, coin: str) -> int:
        """Hyperliquid caps max leverage per asset (BTC/ETH 50×, SOL 20×,
        memes typically 3-5×). Querying meta() returns each asset's
        `maxLeverage`. Returns a conservative 3 if lookup fails (memes default).

        Results are cached for _MAX_LEV_TTL seconds (same TTL as szDecimals).
        meta() returns the full universe in one call so we populate all coins
        on a cache miss to amortise the cost across concurrent opens.
        """
        import time as _time
        now = _time.time()
        cached = _MAX_LEV_CACHE.get(coin)
        if cached is not None and cached[1] > now:
            return cached[0]
        try:
            meta = await self._run(self._info.meta)
            expiry = now + _MAX_LEV_TTL
            for u in meta.get("universe", []):
                name = u.get("name")
                if name:
                    _MAX_LEV_CACHE[name] = (max(1, int(u.get("maxLeverage", 3))), expiry)
            if coin in _MAX_LEV_CACHE:
                return _MAX_LEV_CACHE[coin][0]
        except Exception as exc:
            logger.warning("_get_max_leverage failed for %s: %s", coin, exc)
        return 3   # safe fallback for unknown / illiquid coin

    async def _clip_leverage(self, coin: str, requested: int) -> int:
        """Cap user's requested leverage at the asset's HL max. Without
        this, opening a 30× position on a meme (which caps at 3×) gets
        rejected by HL and the trade is silently dropped."""
        max_lev = await self._get_max_leverage(coin)
        if requested > max_lev:
            logger.info("Clipped leverage for %s: %d× → %d× (HL max)",
                        coin, requested, max_lev)
            return max_lev
        return requested

    async def _mid_price(self, coin: str) -> float:
        mids = await self._run(self._info.all_mids)
        price = float(mids.get(coin, 0))
        if price <= 0:
            raise ValueError(f"Cannot get mid price for {coin}")
        return price

    # ── public interface ─────────────────────────────────────────────────────

    async def get_balance(self) -> float:
        """Return withdrawable USDC balance of the main (MetaMask) account."""
        try:
            state = await self._run(self._info.user_state, self._account_address)
            return float(state.get("withdrawable", 0))
        except Exception as exc:
            logger.error("get_balance error: %s", exc)
            return 0.0

    async def get_open_positions(self) -> list:
        """Return all open positions of the main account (non-zero size only)."""
        try:
            state = await self._run(self._info.user_state, self._account_address)
            positions = []
            for asset_pos in state.get("assetPositions", []):
                pos = asset_pos.get("position", {})
                szi = float(pos.get("szi", 0))
                if szi != 0:
                    positions.append({
                        "coin": pos.get("coin"),
                        "szi": szi,                                    # positive = long, negative = short
                        "entry_px": float(pos.get("entryPx", 0) or 0),
                        "unrealized_pnl": float(pos.get("unrealizedPnl", 0) or 0),
                        "leverage": pos.get("leverage", {}),
                    })
            return positions
        except Exception as exc:
            logger.error("get_open_positions error: %s", exc)
            raise

    async def set_leverage(self, coin: str) -> int:
        """Set isolated leverage for the given coin, clipped to the asset's
        HL maximum. Returns the actual leverage used (may be less than
        self._leverage if asset capped). Caller should prefer this return
        value when computing notional / position size."""
        effective = await self._clip_leverage(coin, self._leverage)
        try:
            await self._run(
                self._exchange.update_leverage,
                effective,
                coin,
                False,   # is_cross=False → isolated margin
            )
            logger.info("Set leverage %dx for %s (isolated)", effective, coin)
        except Exception as exc:
            # DO NOT swallow — if leverage isn't set, the position would open
            # at whatever HL had previously. Downstream stop math
            # (sys2_protective_stop_pct) uses the return value; a wrong
            # value puts the stop outside the real liquidation line.
            logger.error("set_leverage FAILED for %s %dx: %s — aborting open", coin, effective, exc)
            raise RuntimeError(f"set_leverage failed for {coin} at {effective}×: {exc}") from exc
        return effective

    async def open_position(
        self,
        asset: str,
        side: str,          # "long" or "short"
        size_usd: float,
    ) -> dict:
        """
        Open a market position.

        Returns:
            {"order_id": str, "fill_price": float, "size_coins": float}
        """
        coin = asset.upper()
        is_buy = side.lower() == "long"

        # 1. Set leverage before opening. Capture HL's clipped value — for
        #    illiquid / meme assets the requested leverage may exceed the
        #    asset's max and HL silently clips. Caller MUST use this
        #    `effective_leverage` for downstream risk math (protective stop,
        #    liquidation distance) or stops will fire far short of the real
        #    HL liquidation line.
        effective_leverage = await self.set_leverage(coin)

        # 2. Compute coin size from USD notional
        mid = await self._mid_price(coin)
        sz_dec = await self._get_sz_decimals(coin)
        size_coins = round(size_usd / mid, sz_dec)
        if size_coins <= 0:
            raise ValueError(f"Computed size too small: {size_coins} {coin} from ${size_usd}")

        # 3. Limit price with 1% slippage (IOC acts as market)
        slippage = 0.01
        # Hyperliquid px must be ≤5 sig-figs AND divisible by tick size.
        # For BTC ~$76k tick=$1 (integer); for ETH ~$3k tick=$0.1. Use 5 sig-figs + coin-based rounding.
        raw_px = mid * (1 + slippage) if is_buy else mid * (1 - slippage)
        if raw_px >= 10000:
            limit_px = float(round(raw_px))          # integer dollars for BTC
        elif raw_px >= 1000:
            limit_px = round(raw_px, 1)
        elif raw_px >= 100:
            limit_px = round(raw_px, 2)
        else:
            limit_px = round(raw_px, 3)

        logger.info(
            "Opening %s %s: %.5f coins @ limit %.2f (mid=%.2f, usd=%.2f)",
            side, coin, size_coins, limit_px, mid, size_usd,
        )

        result = await self._run(
            self._exchange.order,
            coin,
            is_buy,
            size_coins,
            limit_px,
            {"limit": {"tif": "Ioc"}},   # Immediate-or-Cancel ≈ market
        )

        statuses = result.get("response", {}).get("data", {}).get("statuses", [{}])
        status = statuses[0] if statuses else {}

        # "error" key present → HL rejected the order entirely
        if "error" in status:
            raise ValueError(f"HL order rejected: {status['error']}")

        filled_info = status.get("filled") or {}
        total_sz = float(filled_info.get("totalSz", 0) or 0)
        if total_sz <= 0:
            # IOC returned with zero fill — no position was opened
            raise ValueError(
                f"IOC order for {coin} returned 0 fill "
                f"(status={status}). No position opened."
            )

        fill_price = float(filled_info.get("avgPx", mid) or mid)

        # oid lives under "filled" for IOC; "resting" is fallback for unexpected GTC
        oid_src = filled_info or status.get("resting") or {}
        order_id = str(oid_src.get("oid", "")) if oid_src else ""

        logger.info("Opened %s %s @ %.2f size=%.5f (order_id=%s, lev=%dx)",
                    side, coin, fill_price, total_sz, order_id, effective_leverage)
        return {
            "order_id": order_id,
            "fill_price": fill_price,
            "size_coins": total_sz,
            "effective_leverage": effective_leverage,
        }

    async def close_position(self, asset: str) -> dict:
        """
        Close all open positions for the given asset using a reduce-only IOC order.

        Returns:
            {"fill_price": float}
        """
        coin = asset.upper()

        positions = await self.get_open_positions()
        target = next((p for p in positions if p.get("coin") == coin), None)
        if target is None:
            logger.warning("No open %s position found on HL — already closed externally?", coin)
            # Return sentinel so callers can detect "nothing to close" vs "closed @ price"
            return {"fill_price": None, "already_closed": True}

        szi = float(target["szi"])
        is_buy = szi < 0          # closing a short → buy back
        size = abs(szi)

        mid = await self._mid_price(coin)
        slippage = 0.01
        # Hyperliquid px must be ≤5 sig-figs AND divisible by tick size.
        # For BTC ~$76k tick=$1 (integer); for ETH ~$3k tick=$0.1. Use 5 sig-figs + coin-based rounding.
        raw_px = mid * (1 + slippage) if is_buy else mid * (1 - slippage)
        if raw_px >= 10000:
            limit_px = float(round(raw_px))          # integer dollars for BTC
        elif raw_px >= 1000:
            limit_px = round(raw_px, 1)
        elif raw_px >= 100:
            limit_px = round(raw_px, 2)
        else:
            limit_px = round(raw_px, 3)

        logger.info("Closing %s %s: size=%.5f @ limit %.2f", coin, "short" if is_buy else "long", size, limit_px)

        result = await self._run(
            self._exchange.order,
            coin,
            is_buy,
            size,
            limit_px,
            {"limit": {"tif": "Ioc"}},
            reduce_only=True,
        )

        statuses = result.get("response", {}).get("data", {}).get("statuses", [{}])
        status = statuses[0] if statuses else {}
        filled_info = status.get("filled") or {}
        total_sz = float(filled_info.get("totalSz", 0) or 0)

        if total_sz <= 0:
            # IOC close got 0 fill — position may have already been closed externally.
            # Re-check to be sure before giving up.
            positions_after = await self.get_open_positions()
            still_open = next((p for p in positions_after if p.get("coin") == coin), None)
            if still_open is None:
                logger.info("Close IOC got 0 fill but position is gone — treated as closed", )
                return {"fill_price": mid, "already_closed": False}
            logger.error("Close IOC got 0 fill and position still open for %s", coin)
            raise ValueError(f"Failed to close {coin} position: IOC returned 0 fill")

        fill_price = float(filled_info.get("avgPx", mid) or mid)
        logger.info("Closed %s position @ %.2f (size=%.5f)", coin, fill_price, total_sz)
        return {"fill_price": fill_price, "already_closed": False}

    async def reduce_position(self, asset: str, fraction: float) -> dict:
        """Partially close `fraction` (0<f<1) of the CURRENT open position
        for `asset` with a reduce-only IOC order. Used by System-2 staged
        de-risk. Returns {"fill_price": float, "closed_fraction": float,
        "already_closed": bool}.

        `closed_fraction` is the fraction of the position that was actually
        closed (filled_size / pre-existing size) so the caller's PnL +
        remaining bookkeeping stays exact even on partial fills.
        """
        coin = asset.upper()
        f = max(0.0, min(1.0, float(fraction)))
        if f <= 0.0:
            return {"fill_price": None, "closed_fraction": 0.0, "already_closed": False}
        if f >= 1.0:
            r = await self.close_position(asset)
            r["closed_fraction"] = 0.0 if r.get("already_closed") else 1.0
            return r

        positions = await self.get_open_positions()
        target = next((p for p in positions if p.get("coin") == coin), None)
        if target is None:
            logger.warning("reduce_position: no open %s on HL — already closed?", coin)
            return {"fill_price": None, "closed_fraction": 0.0, "already_closed": True}

        szi = float(target["szi"])
        is_buy = szi < 0                       # reducing a short → buy back
        pre_size = abs(szi)
        sz_dec = await self._get_sz_decimals(coin)
        size = round(pre_size * f, sz_dec)
        if size <= 0:
            # Fraction rounds to nothing at this asset's size precision —
            # treat as a no-op so the caller can advance the step without
            # an erroneous PnL slice.
            logger.warning("reduce_position: %s fraction %.3f rounds to 0 size (pre=%.6f)",
                           coin, f, pre_size)
            return {"fill_price": None, "closed_fraction": 0.0, "already_closed": False}

        mid = await self._mid_price(coin)
        slippage = 0.01
        raw_px = mid * (1 + slippage) if is_buy else mid * (1 - slippage)
        if raw_px >= 10000:
            limit_px = float(round(raw_px))
        elif raw_px >= 1000:
            limit_px = round(raw_px, 1)
        elif raw_px >= 100:
            limit_px = round(raw_px, 2)
        else:
            limit_px = round(raw_px, 3)

        logger.info("Reducing %s by %.0f%%: size=%.6f @ limit %.2f (pre=%.6f)",
                    coin, f * 100, size, limit_px, pre_size)

        result = await self._run(
            self._exchange.order,
            coin, is_buy, size, limit_px,
            {"limit": {"tif": "Ioc"}},
            reduce_only=True,
        )
        statuses = result.get("response", {}).get("data", {}).get("statuses", [{}])
        status = statuses[0] if statuses else {}
        if "error" in status:
            raise ValueError(f"HL reduce rejected: {status['error']}")
        filled_info = status.get("filled") or {}
        total_sz = float(filled_info.get("totalSz", 0) or 0)
        if total_sz <= 0:
            raise ValueError(f"reduce_position {coin}: IOC returned 0 fill")
        fill_price = float(filled_info.get("avgPx", mid) or mid)
        closed_fraction = (total_sz / pre_size) if pre_size > 0 else 0.0
        logger.info("Reduced %s: closed %.6f / %.6f (%.1f%%) @ %.2f",
                    coin, total_sz, pre_size, closed_fraction * 100, fill_price)
        return {"fill_price": fill_price, "closed_fraction": closed_fraction,
                "already_closed": False}