import asyncio
import logging
import os
import time as pytime
from decimal import Decimal
from dotenv import load_dotenv
from lighter import SignerClient, TransactionApi, ApiClient
import requests
import socket
import traceback
import json
from pydantic import StrictInt
from lighter.models.order_book_details import OrderBookDetails
import time

# --- Logging ---
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

# --- Load env ---
load_dotenv()
PRIVATE_KEY = os.getenv("PRIVATE_KEY")
MARKET = os.getenv("MARKET", "ETH")
BASE_URL = "https://mainnet.zklighter.elliot.ai"
SPREAD = Decimal("0.000625")  # Increased to $2.50 at $4000 ETH (from 0.00025 = $1)
REFRESH_INTERVAL = 10
MARKET_ID = 0
ACCOUNT_INDEX = 89181
API_KEY_INDEX = 3
STOP_LOSS_PNL = Decimal("-0.025")  # Stop-loss at -$0.025 (1.5x avg win of $0.0171)
ATR_THRESHOLD = Decimal("0.9")  # Volatility threshold for ATR
ATR_PERIOD = 14  # Standard 14-period ATR for 1-minute data

# --- Manual Nonce Manager (fallback) ---
class ManualNonceManager:
    def __init__(self, base_url, api_key_index=API_KEY_INDEX, account_index=ACCOUNT_INDEX):
        self.base_url = base_url.rstrip('/')
        self.api_key_index = api_key_index
        self.account_index = account_index
        self.current_nonce = 0
        self.last_fetch_time = 0
        self.fetch_interval = 60

    def get_next_nonce(self):
        current_time = pytime.time()
        if (current_time - self.last_fetch_time) > self.fetch_interval:
            try:
                nonce = self._fetch_nonce_from_api()
                if nonce is not None:
                    self.current_nonce = nonce
                    self.last_fetch_time = current_time
                    logger.info(f"🔄 Fetched fresh nonce: {nonce}")
            except Exception as e:
                logger.warning(f"Failed to fetch fresh nonce: {e}")
        self.current_nonce += 1
        return self.current_nonce

    def _fetch_nonce_from_api(self):
        url = f"{self.base_url}/api/v1/nextNonce"
        params = {'account_index': self.account_index, 'api_key_index': self.api_key_index}
        try:
            response = requests.get(url, params=params, timeout=10)
            logger.debug(f"Nonce response: {response.status_code} - {response.text}")
            if response.status_code == 200:
                data = response.json()
                return int(data['nonce']) if 'nonce' in data else None
            return None
        except Exception:
            return None

# --- Connectivity Check ---
def validate_connectivity(base_url):
    logger.info("=== Validating connectivity ===")
    host = base_url.split('://')[1].split('/')[0]
    try:
        ip = socket.gethostbyname(host)
        logger.info(f"✅ DNS resolved: {host} -> {ip}")
    except socket.gaierror as e:
        logger.error(f"❌ DNS failed for {host}: {e}")
        return False

    test_url = f"{base_url}/api/v1/orderBookDetails?market_id={MARKET_ID}"
    try:
        response = requests.get(test_url, timeout=10)
        logger.info(f"✅ HTTP test: {response.status_code} for {test_url}")
        if response.status_code != 200:
            logger.warning(f"⚠️ Market {MARKET} (id={MARKET_ID}) may not exist or ticker endpoint invalid. Returned: {response.status_code} {response.text[:200]}")
        return True
    except requests.exceptions.RequestException as e:
        logger.error(f"❌ HTTP failed: {e}")
        return False

# --- Test nonce endpoint ---
def test_nonce_endpoint():
    logger.info("=== Testing nonce endpoint ===")
    url = f"{BASE_URL}/api/v1/nextNonce"
    params = {'account_index': ACCOUNT_INDEX, 'api_key_index': API_KEY_INDEX}
    try:
        r = requests.get(url, params=params, timeout=10)
        logger.info(f"Status: {r.status_code}")
        logger.info(f"Response: {r.text[:200]}")
        if r.status_code == 200:
            data = r.json()
            logger.info(f"Nonce data: {data}")
    except Exception as e:
        logger.error(f"Nonce test failed: {e}")

# --- Helpers ---
def fetch_orderbook_details(base_url, market_id):
    url = f"{base_url}/api/v1/orderBookDetails?market_id={market_id}"
    try:
        resp = requests.get(url, timeout=10)
        if resp.status_code == 200:
            return resp.json()
        else:
            logger.warning(f"orderBookDetails HTTP {resp.status_code}: {resp.text[:200]}")
            return None
    except Exception as e:
        logger.error(f"fetch_orderbook_details failed: {e}")
        return None

def fetch_orderbook_parsed(base_url, market_id):
    data = fetch_orderbook_details(base_url, market_id)
    if not data:
        return {"bids": [], "asks": [], "bestBid": "0", "bestAsk": "0", "price_decimals": 6, "size_decimals": 0}
    details = data.get("order_book_details", [{}])[0]
    bids = details.get("bids", [])
    asks = details.get("asks", [])
    if not bids or not asks:
        last_trade_price = details.get("last_trade_price")
        if last_trade_price:
            lp = Decimal(str(last_trade_price))
            bids = [[str(lp * Decimal("0.999999")), "1"]]
            asks = [[str(lp * Decimal("1.000001")), "1"]]
    bestBid = details.get("bestBid") or (bids[0][0] if bids else "0")
    bestAsk = details.get("bestAsk") or (asks[0][0] if asks else "0")
    price_decimals = int(details.get("supported_price_decimals", 6))
    size_decimals = int(details.get("supported_size_decimals", 0))
    return {"bids": bids, "asks": asks, "bestBid": str(bestBid), "bestAsk": str(bestAsk), "price_decimals": price_decimals, "size_decimals": size_decimals}

def fetch_open_orders_rest(base_url, account_index, market_id, auth):
    url = f"{base_url}/api/v1/accountActiveOrders"
    params = {"account_index": account_index, "market_id": market_id}
    try:
        resp = requests.get(url, params=params, timeout=10, headers={"accept": "application/json", "Authorization": auth})
        if resp.status_code == 200:
            data = resp.json()
            orders = data.get("orders") or data.get("open_orders") or (data if isinstance(data, list) else [])
            logger.info(f"Full open orders response: {json.dumps(orders)}")
            return orders
        else:
            logger.warning(f"Open orders fetch failed: {resp.status_code} {resp.text[:200]}")
            return []
    except Exception as e:
        logger.error(f"Error fetching open orders: {e}")
        return []

def calculate_atr(base_url, market_id, period=ATR_PERIOD):
    """Calculate 1-minute ATR based on recent price data."""
    try:
        prices = []
        for _ in range(period):
            ob = fetch_orderbook_parsed(base_url, market_id)
            if not ob.get("bids") or not ob.get("asks"):
                logger.warning("No valid orderbook data for ATR calculation")
                return Decimal("0")
            mid_price = (Decimal(ob["bestBid"]) + Decimal(ob["bestAsk"])) / 2
            prices.append(mid_price)
            time.sleep(60 / period)  # Simulate 1-minute intervals
        if len(prices) < period:
            logger.warning("Insufficient price data for ATR")
            return Decimal("0")
        
        tr_list = []
        for i in range(1, len(prices)):
            high = max(prices[i], prices[i-1])
            low = min(prices[i], prices[i-1])
            tr = high - low
            tr_list.append(tr)
        atr = sum(tr_list) / len(tr_list) if tr_list else Decimal("0")
        logger.info(f"Calculated ATR: {atr}")
        return atr
    except Exception as e:
        logger.error(f"ATR calculation failed: {e}")
        return Decimal("0")

async def get_mid_price(public_api, market_id):
    ob = fetch_orderbook_parsed(BASE_URL, market_id)
    bids = ob.get("bids", [])
    asks = ob.get("asks", [])
    if not bids or not asks:
        raise ValueError("No valid orderbook data")
    best_bid = Decimal(str(bids[0][0]))
    best_ask = Decimal(str(asks[0][0]))
    return (best_bid + best_ask) / 2

# --- Place Orders ---
async def place_orders(signer, api, public_api, use_manual_nonce=False, nonce_manager=None):
    try:
        # Check volatility with ATR
        atr = calculate_atr(BASE_URL, MARKET_ID)
        if atr > ATR_THRESHOLD:
            logger.info(f"Volatility too high (ATR={atr} > {ATR_THRESHOLD}), skipping order placement")
            return

        # Fetch nonce
        if use_manual_nonce and nonce_manager:
            nonce = nonce_manager.get_next_nonce()
        else:
            next_nonce_obj = await api.next_nonce(api_key_index=API_KEY_INDEX, account_index=ACCOUNT_INDEX)
            nonce = int(next_nonce_obj.nonce)

        # Generate auth token
        auth, err = signer.create_auth_token_with_expiry(SignerClient.DEFAULT_10_MIN_AUTH_EXPIRY)
        if err is not None:
            logger.error(f"Auth error: {err}")
            return

        # Fetch open orders to determine market direction
        open_orders = fetch_open_orders_rest(BASE_URL, ACCOUNT_INDEX, MARKET_ID, auth)
        logger.info(f"Open orders found (REST): {len(open_orders)}")

        # Check market direction based on open orders
        has_open_ask = any(order.get("is_ask", 0) == 1 for order in open_orders)
        has_open_bid = any(order.get("is_ask", 0) == 0 for order in open_orders)
        
        # Adjust spreads dynamically
        bid_spread_multiplier = 1
        ask_spread_multiplier = 1
        
        if has_open_ask:
            logger.info("Detected open ask order, market moving down, doubling bid spread")
            bid_spread_multiplier = 2
            if not has_open_bid:
                logger.info("Previous bid likely filled, increasing bid spread further")
                bid_spread_multiplier = 4
        
        if has_open_bid:
            logger.info("Detected open bid order, market moving up, doubling ask spread")
            ask_spread_multiplier = 2
            if not has_open_ask:
                logger.info("Previous ask likely filled, increasing ask spread further")
                ask_spread_multiplier = 4

        # Fetch orderbook and calculate prices
        ob_parsed = fetch_orderbook_parsed(BASE_URL, MARKET_ID)
        best_bid = Decimal(str(ob_parsed["bids"][0][0]))
        best_ask = Decimal(str(ob_parsed["asks"][0][0]))
        mid = (best_bid + best_ask) / 2
        logger.info(f"Mid price: {mid}")

        DESIRED_USD = Decimal("20")
        usd_based_size = (DESIRED_USD / mid).quantize(Decimal("0.0000001"), rounding="ROUND_HALF_UP")
        order_size = max(Decimal("0.005"), usd_based_size)  # Enforce min 0.01 AVAX

        # Adjust spreads based on market direction
        bid_spread = SPREAD * bid_spread_multiplier
        ask_spread = SPREAD * ask_spread_multiplier
        buy_price = (mid * (1 - bid_spread / 2)).quantize(Decimal("0.0000001"))
        sell_price = (mid * (1 + ask_spread / 2)).quantize(Decimal("0.0000001"))
        
        # Calculate stop-loss prices
        stop_loss_buy = (buy_price * (1 + STOP_LOSS_PNL / mid)).quantize(Decimal("0.0000001"))  # For buy: price goes up
        stop_loss_sell = (sell_price * (1 - STOP_LOSS_PNL / mid)).quantize(Decimal("0.0000001"))  # For sell: price goes down
        logger.info(f"Buy price: {buy_price}, Sell price: {sell_price}, Bid spread multiplier: {bid_spread_multiplier}, Ask spread multiplier: {ask_spread_multiplier}")
        logger.info(f"Stop-loss buy: {stop_loss_buy}, Stop-loss sell: {stop_loss_sell}")

        price_decimals = int(ob_parsed.get("price_decimals", 6))
        size_decimals = int(ob_parsed.get("size_decimals", 0))
        scaled_buy_price = int((buy_price * (10 ** price_decimals)).to_integral_value(rounding="ROUND_HALF_UP"))
        scaled_sell_price = int((sell_price * (10 ** price_decimals)).to_integral_value(rounding="ROUND_HALF_UP"))
        scaled_stop_loss_buy = int((stop_loss_buy * (10 ** price_decimals)).to_integral_value(rounding="ROUND_HALF_UP"))
        scaled_stop_loss_sell = int((stop_loss_sell * (10 ** price_decimals)).to_integral_value(rounding="ROUND_HALF_UP"))
        scaled_size = int((order_size * (10 ** size_decimals)).to_integral_value(rounding="ROUND_HALF_UP"))

        base_client_order_index = int(pytime.time() * 1000)
        next_client_idx = 0
        expiration_time_ms = int(time.time() * 1000) + (5 * 60 * 1000)

        # Place BUY limit order
        tx_info_buy, err_buy = signer.sign_create_order(
            int(MARKET_ID),
            int(base_client_order_index + next_client_idx),
            int(scaled_size),
            int(scaled_buy_price),
            int(0),  # is_ask
            int(0),  # LIMIT
            int(SignerClient.ORDER_TIME_IN_FORCE_GOOD_TILL_TIME),
            int(0),  # reduce_only
            int(0),  # trigger_price
            int(SignerClient.DEFAULT_28_DAY_ORDER_EXPIRY),
            int(nonce)
        )
        if err_buy:
            logger.error(f"Sign buy error: {err_buy}")
        else:
            tx_info_buy_str = tx_info_buy if isinstance(tx_info_buy, str) else json.dumps(tx_info_buy)
            try:
                response = await api.send_tx(tx_type=StrictInt(signer.TX_TYPE_CREATE_ORDER), tx_info=tx_info_buy_str)
                logger.info(f"✅ Placed BUY {order_size} at {buy_price}. Send tx response: {response}")
            except Exception as e:
                logger.error(f"Send tx error for BUY: {e}")
            nonce += 1

        next_client_idx += 1

        # Place BUY stop-loss order (sell if price rises to stop_loss_buy)
        tx_info_buy_stop, err_buy_stop = signer.sign_create_order(
            int(MARKET_ID),
            int(base_client_order_index + next_client_idx),
            int(scaled_size),
            int(scaled_stop_loss_buy),
            int(1),  # is_ask (sell to close buy)
            int(2),  # STOP
            int(SignerClient.ORDER_TIME_IN_FORCE_GOOD_TILL_TIME),
            int(0),  # reduce_only
            int(scaled_stop_loss_buy),  # trigger_price
            int(SignerClient.DEFAULT_28_DAY_ORDER_EXPIRY),
            int(nonce)
        )
        if err_buy_stop:
            logger.error(f"Sign buy stop-loss error: {err_buy_stop}")
        else:
            tx_info_buy_stop_str = tx_info_buy_stop if isinstance(tx_info_buy_stop, str) else json.dumps(tx_info_buy_stop)
            try:
                response = await api.send_tx(tx_type=StrictInt(signer.TX_TYPE_CREATE_ORDER), tx_info=tx_info_buy_stop_str)
                logger.info(f"✅ Placed BUY stop-loss {order_size} at {stop_loss_buy}. Send tx response: {response}")
            except Exception as e:
                logger.error(f"Send tx error for BUY stop-loss: {e}")
            nonce += 1

        next_client_idx += 1

        # Place SELL limit order
        tx_info_sell, err_sell = signer.sign_create_order(
            int(MARKET_ID),
            int(base_client_order_index + next_client_idx),
            int(scaled_size),
            int(scaled_sell_price),
            int(1),  # is_ask
            int(0),  # LIMIT
            int(SignerClient.ORDER_TIME_IN_FORCE_GOOD_TILL_TIME),
            int(0),  # reduce_only
            int(0),  # trigger_price
            int(SignerClient.DEFAULT_28_DAY_ORDER_EXPIRY),
            int(nonce)
        )
        if err_sell:
            logger.error(f"Sign sell error: {err_sell}")
        else:
            tx_info_sell_str = tx_info_sell if isinstance(tx_info_sell, str) else json.dumps(tx_info_sell)
            try:
                response = await api.send_tx(tx_type=StrictInt(signer.TX_TYPE_CREATE_ORDER), tx_info=tx_info_sell_str)
                logger.info(f"✅ Placed SELL {order_size} at {sell_price}. Send tx response: {response}")
            except Exception as e:
                logger.error(f"Send tx error for SELL: {e}")
            nonce += 1

        next_client_idx += 1

        # Place SELL stop-loss order (buy if price falls to stop_loss_sell)
        tx_info_sell_stop, err_sell_stop = signer.sign_create_order(
            int(MARKET_ID),
            int(base_client_order_index + next_client_idx),
            int(scaled_size),
            int(scaled_stop_loss_sell),
            int(0),  # is_ask (buy to close sell)
            int(2),  # STOP
            int(SignerClient.ORDER_TIME_IN_FORCE_GOOD_TILL_TIME),
            int(0),  # reduce_only
            int(scaled_stop_loss_sell),  # trigger_price
            int(SignerClient.DEFAULT_28_DAY_ORDER_EXPIRY),
            int(nonce)
        )
        if err_sell_stop:
            logger.error(f"Sign sell stop-loss error: {err_sell_stop}")
        else:
            tx_info_sell_stop_str = tx_info_sell_stop if isinstance(tx_info_sell_stop, str) else json.dumps(tx_info_sell_stop)
            try:
                response = await api.send_tx(tx_type=StrictInt(signer.TX_TYPE_CREATE_ORDER), tx_info=tx_info_sell_stop_str)
                logger.info(f"✅ Placed SELL stop-loss {order_size} at {stop_loss_sell}. Send tx response: {response}")
            except Exception as e:
                logger.error(f"Send tx error for SELL stop-loss: {e}")
            nonce += 1

        # Log open orders to verify new orders were added
        open_orders_after_place = fetch_open_orders_rest(BASE_URL, ACCOUNT_INDEX, MARKET_ID, auth)
        logger.info(f"Open orders after placement: {len(open_orders_after_place)}")

    except Exception as e:
        logger.error(f"Error in MM loop: {e}")
        logger.error(traceback.format_exc())

# --- Main ---
async def main():
    if not PRIVATE_KEY:
        logger.error("PRIVATE_KEY not set")
        return

    use_manual_nonce = False
    nonce_manager = None
    if use_manual_nonce:
        nonce_manager = ManualNonceManager(BASE_URL)

    try:
        signer = SignerClient(private_key=PRIVATE_KEY, url=BASE_URL, api_key_index=API_KEY_INDEX, account_index=ACCOUNT_INDEX)
        api = TransactionApi()
        public_api = ApiClient()
        logger.info(f"Starting MM bot on {MARKET} (id={MARKET_ID}) with {BASE_URL}")
    except Exception as e:
        logger.error(f"Client init failed: {e}. Falling back to manual nonce.")
        use_manual_nonce = True
        nonce_manager = ManualNonceManager(BASE_URL)
        signer = SignerClient(private_key=PRIVATE_KEY, url=BASE_URL, api_key_index=API_KEY_INDEX, account_index=ACCOUNT_INDEX)
        api = TransactionApi()
        public_api = ApiClient()

    while True:
        await place_orders(signer, api, public_api, use_manual_nonce, nonce_manager)
        await asyncio.sleep(REFRESH_INTERVAL)

if __name__ == "__main__":
    if not validate_connectivity(BASE_URL):
        logger.error("❌ Connectivity validation failed. Exiting.")
        exit(1)

    test_nonce_endpoint()
    asyncio.run(main())