from .ltx_director import LTXDirector
from .ltx_director_guide import LTXDirectorGuide, LTXDirectorCropGuides
from .msr_character import MSRCharacter, MSRCharacterSet
from .timeline_nodes import TimelineSection, TimelineSectionSet
from comfy_api.latest import ComfyExtension, io
from typing_extensions import override
from .latent_slice import CleanLatentSlice

class PromptRelay(ComfyExtension):
    @override
    async def get_node_list(self) -> list[type[io.ComfyNode]]:
        return [
            MSRCharacter,
            MSRCharacterSet,
            TimelineSection,
            TimelineSectionSet,
            LTXDirector,
        ]

async def comfy_entrypoint() -> PromptRelay:
    return PromptRelay()

NODE_CLASS_MAPPINGS = {
    "LTXDirectorCS": LTXDirector,
    "MSRCharacterCS": MSRCharacter,
    "MSRCharacterSetCS": MSRCharacterSet,
    "TimelineSectionCS": TimelineSection,
    "TimelineSectionSetCS": TimelineSectionSet,
    "LTXDirectorGuideCS": LTXDirectorGuide,
    "LTXDirectorCropGuidesCS": LTXDirectorCropGuides,
    "CleanLatentSliceCS": CleanLatentSlice,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "LTXDirectorCS": "LTX Director DUMAS",
    "MSRCharacterCS": "MSR Character DUMAS",
    "MSRCharacterSetCS": "MSR Character Set DUMAS",
    "TimelineSectionCS": "Timeline Section DUMAS",
    "TimelineSectionSetCS": "Timeline Section Set DUMAS",
    "LTXDirectorGuideCS": "LTX Director Guide DUMAS",
    "LTXDirectorCropGuidesCS": "LTX Director Crop Guides DUMAS",
    "CleanLatentSliceCS": "Clean Latent Slice DUMAS",
}

WEB_DIRECTORY = "./js"

__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY']