import uuid
from typing import Dict, Literal, Protocol, List

from blackd.middlewares import middleware
from opentrons.protocol_api import Labware, ProtocolContext, InstrumentContext, TrashBin, WasteChute, Well
from opentrons.protocol_api.protocol_context import ModuleTypes
from opentrons.types import DeckSlotName, Location

from dbo.instrument import InstrumentWrapper, InstrumentData
from dbo.labware import LabwareWrapper, LabwareData
from dbo.module import ModuleWrapper, ModuleData
from dbo.trashbin import TrashBinWrapper, TrashBinData

fixed_trash_id = uuid.UUID("8c1004ae-a3b0-4230-a3d1-cf5043d7088c")
waste_chute_id = uuid.UUID("e4114551-24cb-4e5e-9068-38422be95b5b")


class Data:
    def __init__(self, protocol: ProtocolContext, is_ot_flex=True):
        self.protocol: ProtocolContext = protocol
        self.static_items: StaticItemsDict = StaticItemsDict(self)
        self.dynamic_items: DynamicItemsDict = DynamicItemsDict(self)
        self.is_ot_flex: bool = is_ot_flex

    def get_well_location_trash_by_id(self, item_id: uuid.UUID):
        location: Well | Location | TrashBin | WasteChute | None = self.dynamic_items.get_well_location_by_id(item_id)
        trash: TrashBinWrapper | WasteChute | None = None

        if location is None:
            trash = self.static_items.get_trash_by_id(item_id)

        if trash is None:
            return None

        if isinstance(trash, TrashBinWrapper):
            location = trash.trash_bin_instance
        else:
            location = trash

        return location

    def generate_new_id(self):
        generated_uuid = uuid.uuid4()
        while (self.static_items.dict.get(generated_uuid) is not None
               or self.dynamic_items.dict.get(generated_uuid) is not None
               or fixed_trash_id == generated_uuid
               or waste_chute_id == generated_uuid):
            generated_uuid = uuid.uuid4()
        return generated_uuid


class StaticItemsDict:
    dict: Dict[uuid.UUID, LabwareWrapper | ModuleWrapper | InstrumentWrapper | TrashBinWrapper | WasteChute] = {}
    data: Data

    def __init__(self, data: Data):
        self.has_fixed_trash: bool = False
        self.has_waste_chute: bool = False
        self.data = data

    # def apply(self, protocol: ProtocolContext):
    #     if self.has_fixed_trash:
    #         self.register_trash(self.get_fixed_trash_id(), protocol.fixed_trash)
    #
    #     if self.has_waste_chute:
    #        self.register_trash(self.get_waste_chute_id(), protocol.load_waste_chute())
    #
    #     for key, value in self.dict.items():
    #         match value:
    #             case WasteChuteWrapper():
    #                 value.
    #
    #
    #         if isinstance(value, WasteChuteWrapper):
    #             value.waste_chute_instance = protocol.load_waste_chute()
    #         else:
    #             if isinstance(value, TrashBinWrapper):
    #                 if value.data.is_fixed:
    #                     value.trash_bin_instance = protocol.fixed_trash
    #                 else:
    #                     value.trash_bin_instance = protocol.load_trash_bin(str(value.data.deck_location))
    #             else:
    #                 if isinstance(value, ModuleWrapper):
    #                     value.module_instance
    #
    #
    #
    #     """Applies all the objects to the protocol given"""

    def get_generic_by_id(self, item_id: uuid.UUID):
        return self.dict.get(item_id)

    def get_generic_id_by_reference(self, reference: Labware | ModuleTypes | InstrumentContext | TrashBin | WasteChute | None):
        if reference == None:
            return None
        item = None
        for key, value in self.dict.items():
            if value == reference:
                item = key
        return item

    def update_generic_by_id(self, item_id: uuid.UUID):
        item = self.get_generic_by_id(item_id)
        self.update_generic(item)

    def update_generic(self, item: LabwareWrapper | ModuleWrapper | InstrumentWrapper | TrashBinWrapper | WasteChute):
        if isinstance(item, LabwareWrapper):
            item.data.update(item.labware_instance, self)
            return
        if isinstance(item, ModuleWrapper):
            item.data.update(item.module_instance, self)
            return
        if isinstance(item, InstrumentWrapper):
            item.data.update(item.instrument_instance, self)
            return

    def update_all(self):
        for key, item in self.dict.items():
            self.update_generic(item)

    def register_labware(self, labware: Labware):
        wrapper = LabwareWrapper.from_instance(labware, self)
        child = wrapper.labware_instance.child
        if child is not None:
            self.register_labware(child)

        new_uuid = self.generate_new_id()
        self.dict[new_uuid] = wrapper
        return new_uuid

    def unregister_labware(self, item_id: uuid.UUID):
        labware = self.get_labware_by_id(item_id)

        to_delete = [item_id]

        if labware is None:
            return

        if labware is not None:
            child = labware.data.child
            while child is not None:
                to_delete.append(child)
                child_instance = self.get_labware_by_id(child)
                child = child_instance.data.child

        for to_delete_id in to_delete:
            labware = self.get_labware_by_id(to_delete_id)
            if labware is not None:
                self.data.dynamic_items.unregister_wells(labware.labware_instance)
            self.dict.pop(to_delete_id)

    def get_labware_by_id(self, item_id: uuid.UUID):
        item = self.dict.get(item_id)

        if isinstance(item, LabwareWrapper):
            return item

        return None

    def get_labware_id_by_reference(self, labware: Labware):
        item = None
        for key, value in self.dict.items():
            if value == labware:
                return key
        return item

    def get_labware_by_deck_location(self, location: str):
        items = {}

        for key, value in self.dict.items():
            if isinstance(value, LabwareWrapper) and value.data.location == location:
                items[key] = value

        return items

    def register_module(self, module: ModuleTypes):
        wrapper = ModuleWrapper.from_instance(module, self)
        new_uuid = self.generate_new_id()
        self.dict[new_uuid] = wrapper
        return new_uuid

    def get_module_by_id(self, item_id: uuid.UUID):
        item = self.dict.get(item_id)

        if isinstance(item, ModuleWrapper):
            return item

        return None

    def get_module_id_by_reference(self, module: ModuleTypes):
        item = None
        for key, value in self.dict.items():
            if value == module:
                item = key

        return item

    def get_module_by_deck_location(self, location: str):
        item = None
        for key, value in self.dict.items():
            if isinstance(value, ModuleWrapper) and value.data.location == location:
                item = value

        return item

    def register_instrument(self, instrument: InstrumentContext):
        wrapper = InstrumentWrapper.from_instance(instrument, self)
        new_uuid = self.generate_new_id()
        self.dict[new_uuid] = wrapper
        return new_uuid

    def get_instrument_by_id(self, item_id: uuid.UUID):
        item = self.dict.get(item_id)

        if isinstance(item, InstrumentWrapper):
            return item

        return None

    def get_instrument_by_mount(self, mount: str):
        item = None
        for key, value in self.dict.items():
            if isinstance(value, InstrumentWrapper) and value.data.mount == mount:
                item = value

        return item

    def get_instrument_id_by_reference(self, instance: InstrumentContext):
        item = None
        for key, value in self.dict.items():
            if isinstance(value, InstrumentWrapper) and value.instrument_instance == instance:
                item = key

        return item

    def register_trash(self, trash: TrashBin | WasteChute, is_fixed: bool = False, is_waste_chute: bool = False, location: str = None):
        if is_fixed and isinstance(trash, TrashBin):
            self.has_fixed_trash = True
            self.dict[fixed_trash_id] = TrashBinWrapper(trash_bin_instance=trash, data=TrashBinData(deck_location=location, is_fixed=is_fixed))
            return fixed_trash_id

        if is_waste_chute and isinstance(trash, WasteChute):
            self.has_waste_chute = True
            self.dict[waste_chute_id] = trash
            return waste_chute_id

        if isinstance(trash, TrashBin):
            new_uuid = self.generate_new_id()
            self.dict[new_uuid] = TrashBinWrapper(trash_bin_instance=trash, data=TrashBinData(deck_location=location, is_fixed=False))
            return new_uuid

    def get_trash_by_id(self, item_id: uuid.UUID):
        item = self.dict.get(item_id)

        if isinstance(item, TrashBinWrapper):
            return item

        if isinstance(item, WasteChute):
            return item

        return None

    def get_trash_id_by_deck_location(self, location: str):
        item = None

        for key, value in self.dict.items():
            if isinstance(value, TrashBinWrapper) and value.data.deck_location == location:
                item = key

        return item

    def generate_new_id(self):
        return self.data.generate_new_id()


class DynamicItemsDict:
    dict: Dict[uuid.UUID, Well | Location] = {}
    data: Data

    def __init__(self, data: Data):
        self.data = data

    def get_well_location_by_id(self, item_id: uuid.UUID):
        return self.dict.get(item_id)

    def register_location(self, location: Location):
        for key, value in self.dict.items():
            if value == location:
                return key

        new_uuid = self.generate_new_id()
        self.dict[new_uuid] = location
        return new_uuid

    def unregister_locations(self, labware: Labware):
        for key, value in self.dict.items():
            if isinstance(value, Location) and value.labware == labware:
                self.dict.pop(key)

    def get_location_by_id(self, item_id: uuid.UUID):
        item = self.dict.get(item_id)

        if isinstance(item, Location):
            return item

        return None

    def register_wells_by_labware(self, labware: Labware):
        for key, value in self.dict.items():
            if value.parent == labware:
                return

        for well in labware.wells():
            new_uuid = self.generate_new_id()
            self.dict[new_uuid] = well

    def unregister_wells(self, labware: Labware):
        for well in labware.wells():
            if well.parent == labware:
                self.dict.pop(self.get_well_id_by_reference(well))

    def get_well_by_id(self, item_id: uuid.UUID):
        item = self.dict.get(item_id)

        if isinstance(item, Well):
            return item

        return None

    def get_well_id_by_reference(self, well: Well):
        item = None
        for key, value in self.dict.items():
            if value == well:
                return key
        return item

    def get_well_ids_by_reference(self, wells: List[Well]):
        items = []
        for well in wells:
            for key, value in self.dict.items():
                if value == well:
                    items.append(key)
        return items

    def generate_new_id(self):
        return self.data.generate_new_id()
