from datetime import datetime, timezone
from dotenv import load_dotenv
from ibkr import Client
from option_type import CALL, PUT
from os import getenv
from tastytrade import Tastytrade
from tastytrade.order import create_credit_spread, create_stop_limit_order
from tastytrade.symbology import zero_dte_spx_contract as contract
from time import sleep

load_dotenv()

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

tastytrade_account = getenv('TASTYTRADE_ACCOUNT')
tastytrade_username = getenv('TASTYTRADE_USERNAME')
tastytrade_password = getenv('TASTYTRADE_PASSWORD')
tastytrade_client = Tastytrade(tastytrade_username, tastytrade_password)
tastytrade_client.login()

underlying_ticker = ibkr_client.get_ticker('SPX', 'CBOE')
current_price = underlying_ticker.last

# Filtering strikes based on distance from current price speeds up the request.
max_strike_distance = 100
def contract_filter(contract):
    if contract.right == 'C':
        return contract.strike <= (current_price + max_strike_distance) and contract.strike >= current_price
    return contract.strike <= current_price and contract.strike >= (current_price - max_strike_distance)

# The weekly symbol for SPX (SPXW) is required in order to distinguish from monthly options.
option_chain = ibkr_client.get_option_chain('SPX', datetime.now(), sub_symbol = 'SPXW', contract_filter = contract_filter)
print(option_chain)

target_delta = 0.20

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'])

put_spread_limit_price = short_put_contract['Bid'] - long_put_contract['Ask'] - 0.05 # Yield to the MMs.
call_spread_limit_price = short_call_contract['Bid'] - long_call_contract['Ask'] - 0.05

print("Short Put Strike:", short_put_strike)
print('Long Put Strike:', long_put_strike)
print('Put Spread Limit Price:', put_spread_limit_price)
print('Short Call Strike:', short_call_strike)
print('Long Call Strike:', long_call_strike)
print('Call Spread Limit Price:', call_spread_limit_price)

put_credit_spread = create_credit_spread(
    contract(PUT, short_put_strike), 
    contract(PUT, long_put_strike), 
    put_spread_limit_price, 1
)
call_credit_spread = create_credit_spread(
    contract(CALL, short_call_strike), 
    contract(CALL, long_call_strike), 
    call_spread_limit_price, 1
)

entry_time = datetime.now(timezone.utc)

put_spread_result = tastytrade_client.submit_order(tastytrade_account, put_credit_spread)
call_spread_result = tastytrade_client.submit_order(tastytrade_account, call_credit_spread)

print(put_spread_result)
print(call_spread_result)

def spread_fill_price(short_position, long_position):
    short_leg_price = float(short_position['average-open-price'])
    long_leg_price = float(long_position['average-open-price'])
    return short_leg_price - long_leg_price

def fill_time(position):
    return datetime.strptime(position["created-at"], '%Y-%m-%dT%H:%M:%S.%f%z')

def wait_for_fill():
    while True:
        positions = tastytrade_client.get_positions(tastytrade_account)
        positions = positions.get('data', {}).get('items', []) # TODO: Client should handle this.
        print(positions)
        # Consider only positions created after the order was submitted.
        new_positions = [position for position in positions if fill_time(position) > entry_time]

        if len(new_positions) == 4: # Assuming no other positions, 4 legs in an iron condor.
            short_put = next(p for p in new_positions if str(int(short_put_strike)) in p['symbol'])
            long_put = next(p for p in new_positions if str(int(long_put_strike)) in p['symbol'])
            short_call = next(p for p in new_positions if str(int(short_call_strike)) in p['symbol'])
            long_call = next(p for p in new_positions if str(int(long_call_strike)) in p['symbol'])

            put_spread_fill_price = spread_fill_price(short_put, long_put)
            call_spread_fill_price = spread_fill_price(short_call, long_call)
            
            return put_spread_fill_price, call_spread_fill_price

        # If not all positions are filled, sleep for a few seconds, then retry.
        sleep(3)

put_spread_fill_price, call_spread_fill_price = wait_for_fill()

print(f'Put Spread Fill Price: {put_spread_fill_price}')
print(f'Call Spread Fill Price: {call_spread_fill_price}')

put_spread_stop = put_spread_fill_price * 2.0
put_spread_stop_order = create_stop_limit_order(
    contract(PUT, short_put_strike), 
    contract(PUT, long_put_strike), 
    stop_trigger = put_spread_stop - 0.25, # Allow for slippage.
    limit_price = put_spread_stop, 
    quantity = 1
)

call_spread_stop = call_spread_fill_price * 2.0
call_spread_stop_order = create_stop_limit_order(
    contract(CALL, short_call_strike), 
    contract(CALL, long_call_strike), 
    stop_trigger = call_spread_stop - 0.25, # Allow for slippage.
    limit_price = call_spread_stop, 
    quantity = 1
)

put_spread_stop_result = tastytrade_client.submit_order(tastytrade_account, put_spread_stop_order)
call_spread_stop_result = tastytrade_client.submit_order(tastytrade_account, call_spread_stop_order)

print(put_spread_stop_result)
print(call_spread_stop_result)