from QuantConnect.Algorithm import QCAlgorithm
from QuantConnect.Data.Custom import *
from QuantConnect.Orders import *
from QuantConnect.Securities.Option import OptionPriceModels
from datetime import timedelta, datetime
import csv
import io

class SPXWTradingAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2023, 11, 3)
        self.SetCash(100000)
        spx = self.AddIndex("SPX").Symbol
        option = self.AddIndexOption(spx, "SPXW")
        option.SetFilter(lambda universe: universe.IncludeWeeklys().Strikes(-100, 100).Expiration(timedelta(0), timedelta(2)))

        self.symbol = option.Symbol
        option.PriceModel = OptionPriceModels.CrankNicolsonFD()
        option.EnableGreekApproximation = True
        self.last_row = None

        self.Schedule.On(self.DateRules.EveryDay(self.symbol), self.TimeRules.Every(timedelta(seconds=30)), self.Trade)

    def Trade(self):
        try:
            csv_string = self.Download('https://docs.google.com/spreadsheets/d/1wwadCU8msu6FEUJt1ANoZS2qMO2MWiheARrdm7zaQlM/export?format=csv')
            csv_reader = csv.DictReader(io.StringIO(csv_string))
            new_row = None
            for i, order in enumerate(csv_reader):
                new_row = order

            if new_row is None or new_row == self.last_row:
                return

            self.last_row = new_row
            self.Debug(','.join([v for v in self.last_row.values()]))

            expiry = datetime.strptime(str(self.last_row['TWS Contract Date']), '%Y%m%d')
            optionchain = self.OptionChainProvider.GetOptionContractList(self.symbol, self.Time.date())
            contracts = [i for i in optionchain if i.ID.Date.date() == expiry.date()]
            
            for security in self.Securities.Values:
                self.Debug(str(security.Symbol))
            
            if len(contracts) < 1:
                self.Debug(f"Not enough option contracts for order: {self.last_row}")
                return

            legs = []
            total_limit_price = 0

            leg_details = []
            for i in range(1, 5):
                if self.last_row[f'Strike {i}'].strip():
                    strike = float(self.last_row[f'Strike {i}'].strip())
                    contracts.sort(key=lambda x: abs(x.ID.StrikePrice - strike))
                    if len(contracts) > 0:
                        if self.Securities.ContainsKey(contracts[0]):
                            contract = self.Securities[contracts[0]]
                            if self.last_row[f'Right {i}']:
                                right = OptionRight.Call if self.last_row[f'Right {i}'] == 'C' else OptionRight.Put
                                self.Debug(f"Right {i}: {self.last_row[f'Right {i}']}")
                                limit_price = (contract.AskPrice + contract.BidPrice) / 2
                                action = self.last_row[f'Action {i}']
                                if action == 'BUY':
                                    quantity = int(self.last_row['Order Quantity'])
                                elif action == 'SELL':
                                    quantity = -int(self.last_row['Order Quantity'])
                                else:
                                    self.Debug(f"Invalid action: {action}")
                                    continue
                                legs.append(Leg.Create(contracts[0], quantity, limit_price))
                                total_limit_price += limit_price
                                leg_details.append(f"{right} Strike {strike}")
                                self.Debug(f"Added leg: {contracts[0].Symbol.Value}, {right}, {quantity}")
                        else:
                            self.Debug(f"Contract not found in Securities: {contracts[0]}")
                    else:
                        self.Debug(f"No contracts found for strike {strike}")

            self.Debug(f"Last row detected on the sheet has {len(legs)} legs in it, {', '.join(leg_details)} at the expiration date {expiry.strftime('%Y%m%d')}.")

            if legs:
                ticket = self.ComboLegLimitOrder(legs, total_limit_price)
                self.Debug(f"Symbol: {ticket.Symbol}; Quantity filled: {ticket.QuantityFilled}; Fill price: {ticket.AverageFillPrice}")
            else:
                self.Debug("No valid legs for combo order")
        except Exception as e:
            self.Debug(f"Error in Trade: {str(e)}")

    def OnOrderEvent(self, orderEvent):
        if orderEvent.Status == OrderStatus.Filled:
            self.Debug("Order filled: " + str(orderEvent))