from AlgorithmImports import *

'''
    Class to handle bracket orders
'''
class BracketOrder:
    def __init__(self, algorithm, orderId, symbol, parentTicket, totQuantity, stopPrice, profitPrice):
        self.algorithm = algorithm

        self.parentId = orderId
        self.symbol = symbol
        self.parentTicket = parentTicket
        self.totQuantity = totQuantity
        self.stopPrice = stopPrice
        self.profitPrice = profitPrice

        # Initialize future properties
        self.filledQuantity = 0
        self.stopOrderTicket = None
        self.profitOrderTicket = None
    
    '''
        Handle order changes

        Orderstatus enum:
        https://www.quantconnect.com/docs/algorithm-reference/trading-and-orders#Trading-and-Orders-Tracking-Order-Events
    '''
    def HandleOrderChange(self, orderEvent):
        eventFillQuantity = orderEvent.FillQuantity

        # If the orderstatus is PartiallyFilled or Filled it will have a filledQuantity > 0
        if (eventFillQuantity > 0):
            self.filledQuantity += eventFillQuantity

            # Place/add to stop order (negative because we are selling)
            self.PlaceStopOrder(-eventFillQuantity)

            # Place/add to profit order
            self.PlaceProfitOrder(eventFillQuantity)
    
    '''
        Place/add to stop market order when the parent order is (partially) filled
    '''
    def PlaceStopOrder(self, stopQuantity):
        # Create new profit order ticket if non exists
        if (self.stopOrderTicket is None):
            self.stopOrderTicket = self.algorithm.StopMarketOrder(self.symbol, stopQuantity, self.stopPrice)
            self.algorithm.bracketOrders[self.stopOrderTicket.OrderId] = self
            self.algorithm.Debug(f"AFTER STOP ORDER, {self.algorithm.bracketOrders.keys()}")

        # Update existing profit order ticket
        updateSettings = UpdateOrderFields()
        updateSettings.Quantity = stopQuantity
        response = self.stopOrderTicket.Update(updateSettings)


    '''
        Place/add to profit taking order when the parent order is (partially) filled
    '''
    def PlaceProfitOrder(self, profitQuantity):
        # Create new profit order ticket if non exists
        if (self.profitOrderTicket is None):
            self.profitOrderTicket = self.algorithm.LimitOrder(self.symbol, profitQuantity, self.profitPrice)
            self.algorithm.bracketOrders[self.profitOrderTicket.OrderId] = self

        # Update existing profit order ticket
        updateSettings = UpdateOrderFields()
        updateSettings.Quantity = profitQuantity
        response = self.profitOrderTicket.Update(updateSettings)    

from AlgorithmImports import *
from BracketOrder import BracketOrder

class TestingOrders(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2020, 5, 5)  # Set Start Date
        self.SetEndDate(2020, 5, 10)  # Set End Date
        self.SetCash(100000)  # Set Strategy Cash

        # Add initial assets
        self.openingBarDict = {}
        self.assets = ["GOOG","AMZN", "JNJ", "PEP", "BAC"]
        for asset in self.assets:
            self.AddEquity(asset, Resolution.Minute).SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        # Initiate bracket order handling dictionary
        self.bracketOrders = {}

    def OnData(self, data: Slice):
        """OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        """

        for bar in data.Bars.Values:
            if (bar.Time.hour == 9 and bar.Time.minute == 30):
                ticker = bar.Symbol
                # Create stop limit order on break of max open/close
                placementPrice = max(bar.Close, bar.Open)
                limitPrice = placementPrice + 0.05
                qty = 10
                parentTicket = self.StopLimitOrder(ticker, qty, placementPrice, limitPrice) # Can I await this and the next line
                self.bracketOrders[parentTicket.OrderId] = BracketOrder(self, parentTicket.OrderId, parentTicket.Symbol, parentTicket, qty, placementPrice - 1, placementPrice + 1) # Can I await this
        
    def OnOrderEvent(self, orderEvent: OrderEvent):
        # Check if the order event ID is allready in the bracket order dict
        # If not, create a new BracketOrder object
        self.Debug(f"Length: {len(self.bracketOrders)}")
        #bracketOrder = self.bracketOrders[orderEvent.OrderId] --- CRASHES BECAUSE ORDER TICKET ID IS NOT A KEY YET