from ib_insync import *
from datetime import datetime, timedelta
import time

# Set up Interactive Brokers connection
ib = IB()
ib.connect('127.0.0.1', 7497, clientId=1)  # Modify port and clientId as necessary

# Define the trading parameters
symbol = 'GOOGL'
contract = Stock(symbol, 'SMART', 'USD')

# Define the time for trading
time_after_open = timedelta(minutes=2)
time_before_close = timedelta(minutes=2)

# Helper function to get the current price
def get_current_price(contract):
    ticker = ib.reqMktData(contract, '', False, False)
    ib.sleep(1)
    return (ticker.bid + ticker.ask) / 2

# Helper function to get the current account balance
def get_account_balance():
    account_values = ib.accountSummary()
    cash_value = next(item for item in account_values if item.tag == 'AvailableFunds').value
    return float(cash_value)

# Helper function to place orders
def place_order(contract, quantity, action):
    order = MarketOrder(action, quantity)
    trade = ib.placeOrder(contract, order)
    ib.sleep(1)
    return trade

# Helper function to close all positions
def close_all_positions():
    positions = ib.positions()
    for position in positions:
        contract = position.contract
        quantity = position.position
        action = 'SELL' if quantity > 0 else 'BUY'
        place_order(contract, abs(quantity), action)
    print("Closed all positions.")

# Trading logic
def trade():
    initial_price = get_current_price(contract)
    peak_price = initial_price
    trough_price = initial_price

    while True:
        current_time = datetime.now()

        # Get market open/close times
        trading_hours = ib.reqContractDetails(contract)[0].tradingHours
        market_open_time = datetime.strptime(trading_hours.split(';')[0].split('-')[0], '%Y%m%d:%H%M')
        market_close_time = datetime.strptime(trading_hours.split(';')[0].split('-')[1], '%Y%m%d:%H%M')

        # Calculate trade start and stop times
        start_trading_time = market_open_time + time_after_open
        stop_trading_time = market_close_time - time_before_close

        # Check if current time is within trading hours
        if start_trading_time <= current_time <= stop_trading_time:
            current_price = get_current_price(contract)
            balance = get_account_balance()
            price_change = (current_price - initial_price) / initial_price

            # Update the peak and trough prices
            if current_price > peak_price:
                peak_price = current_price
            if current_price < trough_price:
                trough_price = current_price

            # Buy when there's a gain, using all available balance
            if current_price > initial_price and balance > 1:
                quantity = balance / current_price
                place_order(contract, quantity, 'BUY')
                print(f"Bought ${balance} worth of {symbol} at {current_price} due to price gain.")
                wait_for_balance_to_update()

            # Sell when there's a loss, using the entire position
            if current_price < initial_price and balance < 1:
                position_qty = sum(pos.position for pos in ib.positions() if pos.contract.symbol == symbol)
                place_order(contract, position_qty, 'SELL')
                print(f"Sold {position_qty} shares of {symbol} at {current_price} due to price loss.")
                wait_for_balance_to_update()

            # Short when there's a loss, using all available balance
            if current_price < initial_price and balance > 1:
                quantity = balance / current_price
                place_order(contract, quantity, 'SELL')
                print(f"Shorted ${balance} worth of {symbol} at {current_price} due to price loss.")
                wait_for_balance_to_update()

            # Cover short when there's a gain, using the entire short position
            if current_price > initial_price and balance < 1:
                position_qty = sum(pos.position for pos in ib.positions() if pos.contract.symbol == symbol)
                place_order(contract, position_qty, 'BUY')
                print(f"Covered short {position_qty} shares of {symbol} at {current_price} due to price gain.")
                wait_for_balance_to_update()

        # Close all positions a couple of minutes before market close
        if current_time > stop_trading_time:
            close_all_positions()
            break  # Exit the loop after closing positions for the day

        time.sleep(1)  # Check as frequently as possible

# Helper function to wait until the account balance updates after a sell
def wait_for_balance_to_update():
    while True:
        balance = get_account_balance()
        if balance > 1:  # Wait until there is at least $1 available to trade again
            break
        time.sleep(1)

# Run the trading strategy
trade()

# Disconnect IB API after trading session
ib.disconnect()