#!/usr/bin/env python3
"""
FTDT Quant Lab — Emergency Kill Switch.

Flattens all active positions and cancels all open orders.
Runs against either testnet or mainnet Hyperliquid.

Usage:
    python scripts/kill_switch.py --mainnet              # Production emergency stop
    python scripts/kill_switch.py --testnet              # Testnet safety test
    python scripts/kill_switch.py --testnet --dry-run    # Print what would happen

The script:
  1. Loads private key from environment (HL_PRIVATE_KEY) or keyfile
  2. Fetches current positions from HL API
  3. Issues market/close orders to flatten each position
  4. Cancels all open orders
  5. Verifies positions are zero
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

TESTNET_API = "https://api.hyperliquid-testnet.xyz/info"
MAINNET_API = "https://api.hyperliquid.xyz/info"
TESTNET_EXCHANGE = "https://api.hyperliquid-testnet.xyz/exchange"
MAINNET_EXCHANGE = "https://api.hyperliquid.xyz/exchange"


def get_private_key() -> str:
    pk = os.environ.get("HL_PRIVATE_KEY", "")
    if pk:
        return pk

    keyfile = Path.home() / ".hl" / "key"
    if keyfile.exists():
        return keyfile.read_text().strip()

    return ""


def sign_l1_action(action: dict, private_key: str) -> dict:
    try:
        from eth_account import Account
        import eth_account.messages
    except ImportError:
        print("eth_account not installed. Install with: pip install eth-account")
        sys.exit(1)

    account = Account.from_key(private_key)
    from datetime import datetime, timezone

    typed_data = {
        "domain": {"chainId": 1337},
        "types": {
            "HyperliquidTransaction": [
                {"name": "txType", "type": "string"},
                {"name": "nonce", "type": "uint64"},
            ],
        },
        "primaryType": "HyperliquidTransaction",
        "message": {
            "txType": action.get("type", "order"),
            "nonce": int(time.time() * 1000),
        },
    }

    signable = eth_account.messages.encode_typed_data(
        typed_data["domain"],
        typed_data["types"],
        typed_data["message"],
    )
    sig = account.sign_message(signable)
    action["signature"] = sig.signature.hex()
    return action


def fetch_positions(api_url: str) -> list[dict]:
    import requests
    pk = get_private_key()
    if not pk:
        return []

    try:
        from eth_account import Account
        account = Account.from_key(pk)
        resp = requests.post(
            api_url,
            json={"type": "clearinghouseState", "user": account.address},
            timeout=10,
        )
        data = resp.json()
        positions = []
        if isinstance(data, dict):
            asset_positions = data.get("assetPositions", [])
            for pos in asset_positions:
                p = pos.get("position", {})
                coin = p.get("coin", "?")
                szi = float(p.get("szi", 0))
                if szi != 0:
                    positions.append({
                        "coin": coin,
                        "szi": szi,
                        "entry_px": float(p.get("entryPx", 0)),
                        "leverage": p.get("leverage", {}),
                    })
        return positions
    except Exception as e:
        print(f"Failed to fetch positions: {e}")
        return []


def close_position(exchange_url: str, coin: str, current_szi: float, dry_run: bool) -> bool:
    import requests

    side = "A" if current_szi < 0 else "B"
    close_size = abs(current_szi)

    action = {
        "type": "order",
        "orders": [{
            "a": 0,
            "b": open_price(exchange_url.replace("exchange", "info"), coin, side),
            "p": open_price(exchange_url.replace("exchange", "info"), coin, side),
            "s": round(close_size, 6),
            "r": False,
            "t": {"limit": {"tif": "Ioc"}},
        }],
        "grouping": "na",
    }

    if not dry_run:
        pk = get_private_key()
        if pk:
            action = sign_l1_action(action, pk)
        try:
            resp = requests.post(exchange_url, json=action, timeout=10)
            result = resp.json()
            if resp.status_code == 200:
                print(f"  [{coin}] Closed {close_size:.4f} ({side}) — OK")
                return True
            else:
                print(f"  [{coin}] Close failed: {result}")
                return False
        except Exception as e:
            print(f"  [{coin}] Close error: {e}")
            return False
    else:
        print(f"  [{coin}] [DRY RUN] Would close {close_size:.4f} ({side})")
        return True


def open_price(api_url: str, coin: str, side: str) -> float:
    import requests
    try:
        resp = requests.post(api_url, json={"type": "l2Book", "coin": coin}, timeout=5)
        data = resp.json()
        levels = data.get("levels", [])
        if levels and len(levels) >= 2:
            bids = levels[0]
            asks = levels[1]
            if side == "A":
                px = float(bids[0]["px"]) if bids else 0
                return round(px * 0.99, 1)
            else:
                px = float(asks[0]["px"]) if asks else 0
                return round(px * 1.01, 1)
    except Exception:
        pass
    return 0.0


def cancel_all_orders(exchange_url: str, api_url: str, dry_run: bool):
    import requests
    pk = get_private_key()
    if not pk:
        print("  No private key — cannot cancel orders")
        return

    from eth_account import Account
    account = Account.from_key(pk)

    try:
        resp = requests.post(
            api_url,
            json={"type": "openOrders", "user": account.address},
            timeout=10,
        )
        orders = resp.json()
        if not isinstance(orders, list) or len(orders) == 0:
            print("  No open orders found")
            return

        print(f"  Found {len(orders)} open orders")

        if not dry_run:
            cancels = []
            for order in orders:
                cancels.append({
                    "a": order.get("oid", 0),
                    "b": order.get("coin", "BTC"),
                })

            action = {
                "type": "cancel",
                "cancels": cancels,
            }
            action = sign_l1_action(action, pk)

            try:
                resp = requests.post(exchange_url, json=action, timeout=10)
                if resp.status_code == 200:
                    print(f"  Cancelled {len(cancels)} orders — OK")
                else:
                    print(f"  Cancel failed: {resp.json()}")
            except Exception as e:
                print(f"  Cancel error: {e}")
        else:
            print(f"  [DRY RUN] Would cancel {len(orders)} orders")
    except Exception as e:
        print(f"  Failed to fetch orders: {e}")


def verify_flat(api_url: str) -> bool:
    positions = fetch_positions(api_url)
    if not positions:
        print("  No positions — flat")
        return True

    total = sum(abs(p["szi"]) for p in positions)
    if total < 1e-6:
        print("  No positions — flat")
        return True

    print(f"  WARNING: {len(positions)} positions remain (total exposure: {total:.4f})")
    for p in positions:
        print(f"    {p['coin']}: {p['szi']:.4f} @ ${p['entry_px']:.1f}")
    return False


def main():
    p = argparse.ArgumentParser(description="FTDT Quant Lab — Emergency Kill Switch")
    p.add_argument("--testnet", action="store_true", default=True)
    p.add_argument("--mainnet", dest="testnet", action="store_false")
    p.add_argument("--dry-run", action="store_true", help="Print actions without executing")
    p.add_argument("--skip-cancel", action="store_true", help="Skip order cancellation")
    p.add_argument("--retry", type=int, default=3, help="Retry count for position closes")
    args = p.parse_args()

    api_url = MAINNET_API if not args.testnet else TESTNET_API
    exchange_url = MAINNET_EXCHANGE if not args.testnet else TESTNET_EXCHANGE
    env_name = "MAINNET" if not args.testnet else "TESTNET"

    print("=" * 60)
    print(f"  FTDT Quant Lab — KILL SWITCH [{env_name}]")
    print("=" * 60)
    if args.dry_run:
        print("  *** DRY RUN — no orders will be placed ***")
    print()

    pk = get_private_key()
    if not pk:
        print("  ERROR: No private key found.")
        print("  Set HL_PRIVATE_KEY environment variable or create ~/.hl/key")
        if not args.dry_run:
            sys.exit(1)

    print(f"  [1] Fetching positions...")
    positions = fetch_positions(api_url)
    if not positions:
        print("  No positions to close")
    else:
        total_size = sum(abs(p["szi"]) for p in positions)
        print(f"  Found {len(positions)} position(s) (total size: {total_size:.4f})")
        for p in positions:
            direction = "LONG" if p["szi"] > 0 else "SHORT"
            print(f"    {p['coin']}: {direction} {abs(p['szi']):.4f} @ ${p['entry_px']:.1f}")

        print(f"\n  [2] Closing positions...")
        all_closed = True
        for p in positions:
            for attempt in range(args.retry):
                ok = close_position(exchange_url, p["coin"], p["szi"], args.dry_run)
                if ok:
                    break
                if attempt < args.retry - 1:
                    time.sleep(1)
            else:
                print(f"    [{p['coin']}] Failed after {args.retry} attempts")
                all_closed = False

        if not all_closed and not args.dry_run:
            print("\n  WARNING: Some positions could not be closed!")
        else:
            print("  All positions closed")

    if not args.skip_cancel:
        print(f"\n  [3] Cancelling open orders...")
        cancel_all_orders(exchange_url, api_url, args.dry_run)

    print(f"\n  [4] Verification...")
    if not args.dry_run:
        time.sleep(2)
    ok = verify_flat(api_url)
    if ok:
        print("\n  KILL SWITCH COMPLETE — positions flat, orders cancelled")
    else:
        print("\n  KILL SWITCH WARNING — positions may still exist!")


if __name__ == "__main__":
    main()