import logging
import nest_asyncio
import pandas as pd
import traceback

from database.trades import upsert
from dataclasses import replace
from datetime import datetime
from dotenv import load_dotenv
from ibkr import Client, OptionLeg
from ibkr.option_type import CALL, PUT
from ibkr.order_action import BUY, SELL
from os import getenv
from pytz import timezone
from tradestation import TradeStationClient

load_dotenv()

# Allows for starting an event loop even if there's already one running in the current thread.
# Necessary for monitoring spread prices asynchronously while interacting with the IBKR client.
nest_asyncio.apply()

quantity = 1

def monitor_spread_price(short_leg: OptionLeg, long_leg: OptionLeg, stop_price: float, client: Client):
    """""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
        Stop loss orders will not execute if trying to sell back a contract with no bid while paper trading.
        Therefore, the spread price must be monitored and the spread manually exited if the stop price is reached.
        If there is no bid for the long leg, only the short leg will be exited.
    """""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
    stopped_out = False
    market_data = {}  # Stores real-time market data for each leg.

    def on_market_data_update(update_event):
        # Prevent the trade from being exited multiple times if there are other updates queued.
        # This will prevent unintentionally entering new trades.
        nonlocal stopped_out
        if stopped_out:
            return

        # Ensure there is market data for both legs before proceeding.
        if short_leg in market_data and long_leg in market_data:
            short_contract = market_data[short_leg]
            long_contract = market_data[long_leg]

            # If a contract has no bid -1.0 is returned, set it to 0 to avoid negative mid prices.
            mid_price_short = (max(short_contract.bid, 0) + short_contract.ask) / 2
            mid_price_long = (max(long_contract.bid, 0) + long_contract.ask) / 2
            current_spread_price = mid_price_short - mid_price_long

            logging.info(f'Short Contract: {short_leg.strike} {short_leg.option_type}')
            logging.info(f'Long Contract: {long_leg.strike} {long_leg.option_type}')
            logging.info(f'Current Spread Price: {current_spread_price}')
            logging.info(f'Stop Price: {stop_price}')

            if current_spread_price >= stop_price:
                stopped_out = True
                logging.info('Stop price reached or exceeded. Exiting trade.')

                short_leg_exit = replace(short_leg, action = BUY if short_leg.action == SELL else SELL)
                long_leg_exit = replace(long_leg, action = BUY if long_leg.action == SELL else SELL)
                if long_contract.bid > 0:
                    client.submit_combo_option_order([short_leg_exit, long_leg_exit], quantity)
                    logging.info('Whole spread exited.')
                else:
                    client.submit_option_order(short_leg_exit, quantity)
                    logging.info('Short leg only exited.')

                # Unsubscribe from market data updates once the trade has exited.
                for leg in [short_leg, long_leg]:
                    market_data[leg].updateEvent -= on_market_data_update

    for leg in [short_leg, long_leg]:
        option_contract = client.get_option_contract(leg)
        leg_market_data = client.get_market_data(option_contract, streaming=True)
        market_data[leg] = leg_market_data
        leg_market_data.updateEvent += on_market_data_update

def enter_iron_condor():
    logging.basicConfig(
        filename=f'iron_condor_{datetime.now().strftime("%H%M")}.log', 
        level=logging.INFO,
        format='%(asctime)s : %(levelname)s : %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S'
    )
    try:
        _enter_iron_condor()
    except Exception as e:
        logging.error("Error: %s", traceback.format_exc())

def _enter_iron_condor():
    ibkr_host = getenv('IBKR_HOST')
    ibkr_port = getenv('IBKR_PORT')
    ibkr_client = Client(host = ibkr_host, port = ibkr_port)

    # The weekly symbol for SPX (SPXW) is required in order to distinguish from monthly options.
    symbol, sub_symbol = 'SPX', 'SPXW'
    expiration = datetime.now()
    
    tradestation_client = TradeStationClient(getenv('TRADESTATION_REFRESH_TOKEN'))

    option_chain = tradestation_client.get_options_chain('$SPXW.X', expiration)
    logging.info(option_chain)

    target_delta = 0.10

    def closest_contract_by_delta(target_delta, option_chain, option_type):
        options = option_chain[option_chain['Type'] == option_type].copy()
        options['Delta Distance'] = abs(options['Delta'] - target_delta)
        return options.loc[options['Delta Distance'].idxmin()]

    # Find the strikes that minimize the distance to the target delta.
    short_put_contract = closest_contract_by_delta(-target_delta, option_chain, 'PUT')
    short_call_contract = closest_contract_by_delta(target_delta, option_chain, 'CALL')

    # When selecting long strikes, minimize the distance to a 50 point spread.
    # TODO: Select long strike based on preferred price.
    target_long_put_strike = short_put_contract['Strike'] - 50
    target_long_call_strike = short_call_contract['Strike'] + 50

    def closest_contract_by_strike(target_strike, option_chain, option_type):
        options = option_chain[option_chain['Type'] == option_type].copy()
        options['Strike Distance'] = abs(options['Strike'] - target_strike)
        return options.loc[options['Strike Distance'].idxmin()]

    long_put_contract = closest_contract_by_strike(target_long_put_strike, option_chain, 'PUT')
    long_call_contract = closest_contract_by_strike(target_long_call_strike, option_chain, 'CALL')

    # Build the iron condor.
    short_put_strike = float(short_put_contract['Strike'])
    long_put_strike = float(long_put_contract['Strike'])
    short_call_strike = float(short_call_contract['Strike'])
    long_call_strike = float(long_call_contract['Strike'])

    logging.info(f'Short Put Strike: {short_put_strike}')
    logging.info(f'Long Put Strike: {long_put_strike}')
    logging.info(f'Short Call Strike: {short_call_strike}')
    logging.info(f'Long Call Strike: {long_call_strike}')

    trade_records = []
    call_spread_details = {}
    put_spread_details = {}

    short_call_leg = OptionLeg(symbol, expiration, short_call_strike, CALL, SELL, sub_symbol)
    long_call_leg = OptionLeg(symbol, expiration, long_call_strike, CALL, BUY, sub_symbol)

    call_spread_order = ibkr_client.submit_combo_option_order([short_call_leg, long_call_leg], quantity)
    while not call_spread_order.isDone():
        ibkr_client.ib.waitOnUpdate()

    if call_spread_order.orderStatus.status == 'Filled':
        fill_price = abs(call_spread_order.orderStatus.avgFillPrice)
        logging.info(f'Call Spread Fill Price: {fill_price}')
        monitor_spread_price(short_call_leg, long_call_leg, fill_price * 2, ibkr_client)

        call_spread_details = {
            'Legs': [
                {'Action': 'SELL', 'Strike': short_call_strike, 'Type': 'CALL'},
                {'Action': 'BUY', 'Strike': long_call_strike, 'Type': 'CALL'}
            ],
            'Open': fill_price
        }

    short_put_leg = OptionLeg(symbol, expiration, short_put_strike, PUT, SELL, sub_symbol)
    long_put_leg = OptionLeg(symbol, expiration, long_put_strike, PUT, BUY, sub_symbol)

    put_spread_order = ibkr_client.submit_combo_option_order([short_put_leg, long_put_leg], quantity)
    while not put_spread_order.isDone():
        ibkr_client.ib.waitOnUpdate()

    if put_spread_order.orderStatus.status == 'Filled':
        fill_price = abs(put_spread_order.orderStatus.avgFillPrice)
        logging.info(f'Put Spread Fill Price: {fill_price}')
        monitor_spread_price(short_put_leg, long_put_leg, fill_price * 2, ibkr_client)

        put_spread_details = {
            'Legs': [
                {'Action': 'SELL', 'Strike': short_put_strike, 'Type': 'PUT'},
                {'Action': 'BUY', 'Strike': long_put_strike, 'Type': 'PUT'}
            ],
            'Open': fill_price
        }

    now = datetime.now().astimezone(timezone('US/Eastern'))
    trade_records.append({
        'Date': now.date(),
        'Symbol': symbol,
        'Strategy': f'{int(target_delta * 100)} Delta Iron Condor @ {now.strftime("%H:%M:00")}',
        'Entry Time': now,
        'Exit Time': None,
        'Spreads': [call_spread_details, put_spread_details],
        'Profit': None
    })
    upsert(pd.DataFrame(trade_records))

    # TODO: Add a shutdown hook.
    ibkr_client.run_event_loop()