import argparse
import logging
import os
import sys
import json
import time
from datetime import datetime
import multiprocessing

from eth_account import Account
from hyperliquid.exchange import Exchange
from hyperliquid.info import Info
from hyperliquid.utils import constants
from dotenv import load_dotenv

from logging_utils import setup_logging
from trade_log import log_trade

# Load environment variables from a .env file
load_dotenv()

class TradeExecutor:
    """
    Monitors a shared queue for strategy signals and executes trades.
    This script is now a dedicated, event-driven consumer.
    """

    def __init__(self, log_level: str, trade_signal_queue: multiprocessing.Queue, shared_executor_status: dict = None):
        setup_logging(log_level, 'TradeExecutor')

        self.trade_signal_queue = trade_signal_queue

        # Optional Manager.dict() to store live managed positions and other executor status
        self.shared_executor_status = shared_executor_status

        self.vault_address = os.environ.get("MAIN_WALLET_ADDRESS")
        if not self.vault_address:
            logging.error("MAIN_WALLET_ADDRESS not set.")
            # --- FIX: Raise an exception instead of sys.exit() ---
            # This allows the main_app process manager to catch and log the error.
            raise ValueError("MAIN_WALLET_ADDRESS not set in environment.")
        
        # --- FIX: Corrected constant name from MAIN_NET_API_URL to MAINNET_API_URL ---
        self.info = Info(constants.MAINNET_API_URL, skip_ws=True)
        self.exchanges = self._load_agents()
        if not self.exchanges:
            logging.error("No trading agents found in .env file.")
            # --- FIX: Raise an exception instead of sys.exit() ---
            raise ValueError("No trading agents found in .env file. Check AGENT_PRIVATE_KEY or _AGENT_PK vars.")

        self.managed_positions_path = os.path.join("_data", "executor_managed_positions.json")
        self.managed_positions = self._load_managed_positions()
        logging.info(f"TradeExecutor initialized. Agents available: {list(self.exchanges.keys())}")

    def _load_agents(self) -> dict:
        """
        Discovers and initializes agents by scanning for environment variables.
        """
        exchanges = {}
        logging.info("Discovering agents from environment variables...")
        for env_var, private_key in os.environ.items():
            agent_name = None
            if env_var == "AGENT_PRIVATE_KEY":
                agent_name = "default"
            elif env_var.endswith("_AGENT_PK"):
                agent_name = env_var.replace("_AGENT_PK", "").lower()

            if agent_name and private_key:
                try:
                    agent_account = Account.from_key(private_key)
                    exchanges[agent_name] = Exchange(agent_account, constants.MAINNET_API_URL, account_address=self.vault_address)
                    logging.info(f"Initialized agent '{agent_name}' with address: {agent_account.address}")
                except Exception as e:
                    logging.error(f"Failed to initialize agent '{agent_name}': {e}")
        return exchanges

    def _load_managed_positions(self) -> dict:
        """Loads the state of which strategy manages which position."""
        # Prefer shared in-memory state when available
        try:
            if self.shared_executor_status is not None:
                mgr = self.shared_executor_status.get('managed_positions') if isinstance(self.shared_executor_status, dict) else None
                if mgr:
                    logging.info("Loading managed positions from shared executor status.")
                    return dict(mgr)
        except Exception:
            logging.debug("Unable to read managed positions from shared status. Falling back to file.")

        if os.path.exists(self.managed_positions_path):
            try:
                with open(self.managed_positions_path, 'r') as f:
                    logging.info("Loading existing managed positions state from file.")
                    return json.load(f)
            except (IOError, json.JSONDecodeError):
                logging.warning("Could not read managed positions file. Starting fresh.")
        return {}
    
    def _save_managed_positions(self):
        """Saves the current state of managed positions."""
        try:
            if self.shared_executor_status is not None:
                try:
                    # store under a known key
                    self.shared_executor_status['managed_positions'] = dict(self.managed_positions)
                except Exception:
                    # fallback: try direct assignment
                    self.shared_executor_status['managed_positions'] = self.managed_positions
            else:
                with open(self.managed_positions_path, 'w') as f:
                    json.dump(self.managed_positions, f, indent=4)
        except IOError as e:
            logging.error(f"Failed to save managed positions state: {e}")

    def run(self):
        """
        Main execution loop. Blocks and waits for a signal from the queue.
        """
        logging.info("Trade Executor started. Waiting for signals...")
        while True:
            try:
                trade_signal = self.trade_signal_queue.get()
                if not trade_signal:
                    continue

                logging.info(f"Received signal: {trade_signal}")

                # Basic validation and debug information to help trace gaps
                if 'config' not in trade_signal:
                    logging.error(f"Signal missing 'config' key. Ignoring: {trade_signal}")
                    continue
                if 'strategy_name' not in trade_signal:
                    logging.error(f"Signal missing 'strategy_name' key. Ignoring: {trade_signal}")
                    continue
                # Special command handling
                if isinstance(trade_signal, dict) and trade_signal.get('_cmd') == 'CLOSE_ALL':
                    target_agent = trade_signal.get('agent')
                    logging.warning(f"Received CLOSE_ALL command for agent: {target_agent}")
                    if not target_agent:
                        logging.error("CLOSE_ALL command missing 'agent' field. Ignoring.")
                        continue

                    # Iterate managed positions and close those opened by the target agent
                    to_close = [s for s, v in self.managed_positions.items() if v.get('agent') == target_agent]
                    if not to_close:
                        logging.info(f"No managed positions found for agent '{target_agent}'.")
                        continue

                    for sname in to_close:
                        pos = self.managed_positions.get(sname)
                        if not pos:
                            continue
                        coin = pos.get('coin')
                        side = pos.get('side')
                        size = pos.get('size')
                        # Determine is_buy to neutralize the position
                        is_buy = True if side == 'short' else False
                        logging.warning(f"[CLOSE_ALL] Closing {side} position for strategy {sname}, coin {coin}, size {size}")
                        try:
                            # Use the agent's exchange if available
                            exch = self.exchanges.get(target_agent)
                            if exch:
                                exch.market_open(coin, is_buy, size, None, 0.01)
                            else:
                                logging.error(f"Exchange object for agent '{target_agent}' not found. Skipping live close for {sname}.")
                        except Exception as e:
                            logging.error(f"Error closing position for {sname}: {e}")
                        # remove from managed positions regardless to avoid stuck state
                        try:
                            del self.managed_positions[sname]
                        except KeyError:
                            pass

                    self._save_managed_positions()
                    logging.info(f"CLOSE_ALL for agent '{target_agent}' completed.")
                    continue

                name = trade_signal['strategy_name']
                config = trade_signal['config']
                params = config.get('parameters', {})
                coin = trade_signal['coin']
                desired_signal = trade_signal['signal']
                status = trade_signal 

                size = params.get('size')
                if size is None:
                    logging.error(f"[{name}] No 'size' in parameters: {params}. Skipping.")
                    continue
                leverage_long = int(params.get('leverage_long', 2))
                leverage_short = int(params.get('leverage_short', 2))
                current_position = self.managed_positions.get(name)
                
                agent_name = (config.get("agent") or "default").lower()
                exchange_to_use = self.exchanges.get(agent_name)
                if not exchange_to_use:
                    logging.error(f"[{name}] Agent '{agent_name}' not found. Available agents: {list(self.exchanges.keys())}. Skipping trade.")
                    continue
                
                # --- State Machine Logic (now runs instantly on signal) ---
                if desired_signal == "BUY" or desired_signal == "INIT_BUY":
                    if not current_position:
                        logging.warning(f"[{name}] ACTION: Setting leverage to {leverage_long}x and opening LONG for {coin}.")
                        exchange_to_use.update_leverage(leverage_long, coin)
                        exchange_to_use.market_open(coin, True, size, None, 0.01)
                        self.managed_positions[name] = {"coin": coin, "side": "long", "size": size}
                        log_trade(strategy=name, coin=coin, action="OPEN_LONG", price=status.get('signal_price', 0), size=size, signal=desired_signal)
                    elif current_position['side'] == 'short':
                        logging.warning(f"[{name}] ACTION: Closing SHORT and opening LONG for {coin} with {leverage_long}x leverage.")
                        exchange_to_use.update_leverage(leverage_long, coin)
                        # 1. Close the short by buying back (this is a market_open, but is_buy=True)
                        exchange_to_use.market_open(coin, True, current_position['size'], None, 0.01)
                        log_trade(strategy=name, coin=coin, action="CLOSE_SHORT", price=status.get('signal_price', 0), size=current_position['size'], signal=desired_signal)
                        # 2. Open the new long
                        exchange_to_use.market_open(coin, True, size, None, 0.01)
                        self.managed_positions[name] = {"coin": coin, "side": "long", "size": size}
                        log_trade(strategy=name, coin=coin, action="OPEN_LONG", price=status.get('signal_price', 0), size=size, signal=desired_signal)
                    
                elif desired_signal == "SELL" or desired_signal == "INIT_SELL":
                    if not current_position:
                        logging.warning(f"[{name}] ACTION: Setting leverage to {leverage_short}x and opening SHORT for {coin}.")
                        exchange_to_use.update_leverage(leverage_short, coin)
                        exchange_to_use.market_open(coin, False, size, None, 0.01)
                        self.managed_positions[name] = {"coin": coin, "side": "short", "size": size}
                        log_trade(strategy=name, coin=coin, action="OPEN_SHORT", price=status.get('signal_price', 0), size=size, signal=desired_signal)
                    elif current_position['side'] == 'long':
                        logging.warning(f"[{name}] ACTION: Closing LONG and opening SHORT for {coin} with {leverage_short}x leverage.")
                        exchange_to_use.update_leverage(leverage_short, coin)
                        # 1. Close the long by selling
                        exchange_to_use.market_open(coin, False, current_position['size'], None, 0.01)
                        log_trade(strategy=name, coin=coin, action="CLOSE_LONG", price=status.get('signal_price', 0), size=current_position['size'], signal=desired_signal)
                        # 2. Open the new short
                        exchange_to_use.market_open(coin, False, size, None, 0.01)
                        self.managed_positions[name] = {"coin": coin, "side": "short", "size": size}
                        # --- FIX: Corrected typo from 'signal.desired_signal' to 'signal=desired_signal' ---
                        log_trade(strategy=name, coin=coin, action="OPEN_SHORT", price=status.get('signal_price', 0), size=size, signal=desired_signal)

                elif desired_signal == "FLAT":
                    if current_position:
                        logging.warning(f"[{name}] ACTION: Close {current_position['side']} position for {coin}.")
                        is_buy = current_position['side'] == 'short'
                        exchange_to_use.market_open(coin, is_buy, current_position['size'], None, 0.01)
                        del self.managed_positions[name]
                        log_trade(strategy=name, coin=coin, action=f"CLOSE_{current_position['side'].upper()}", price=status.get('signal_price', 0), size=current_position['size'], signal=desired_signal)
                
                self._save_managed_positions()

            except Exception as e:
                logging.error(f"An error occurred in the main executor loop: {e}", exc_info=True)
                time.sleep(1)

# This script is no longer run directly, but is called by main_app.py