from .blackbox import Blackbox
from injector import inject, singleton, Injector
from ..configuration import BlackboxConf

blackboxConf = Injector().get(BlackboxConf)

def model_loader(lazy=blackboxConf.lazyloading):
    def load(init_fun):
        model = None
        if not lazy:
            model = init_fun()
        def inner():
            nonlocal model
            if model is None:
                model = init_fun()
            return model
        return inner
    return load

@model_loader(lazy=blackboxConf.lazyloading)
def text_and_image_loader():
    from .text_and_image import TextAndImage
    return Injector().get(TextAndImage)

@model_loader(lazy=blackboxConf.lazyloading)
def audio_chat_loader():
    from .audio_chat import AudioChat
    return Injector().get(AudioChat)

@model_loader(lazy=blackboxConf.lazyloading)
def sentiment_loader():
    from .sentiment import Sentiment
    return Injector().get(Sentiment)

@model_loader(lazy=blackboxConf.lazyloading)
def g2e_loader():
    from .g2e import G2E
    return Injector().get(G2E)

@model_loader(lazy=blackboxConf.lazyloading)
def asr_loader():
    from .asr import ASR
    return Injector().get(ASR)

@model_loader(lazy=blackboxConf.lazyloading)
def vlms_loader():
    from .vlms import VLMS
    return Injector().get(VLMS)

@model_loader(lazy=blackboxConf.lazyloading)
def melotts_loader():
    from .melotts import MeloTTS
    return Injector().get(MeloTTS)

@model_loader(lazy=blackboxConf.lazyloading)
def tts_loader():
    from .tts import TTS
    return Injector().get(TTS)

@model_loader(lazy=blackboxConf.lazyloading)
def emotion_loader():
    from .emotion import Emotion
    return Injector().get(Emotion)

@model_loader(lazy=blackboxConf.lazyloading)
def fastchat_loader():
    from .fastchat import Fastchat
    return Injector().get(Fastchat)

@model_loader(lazy=blackboxConf.lazyloading)
def chat_loader():
    from .chat import Chat
    return Injector().get(Chat)

@model_loader(lazy=blackboxConf.lazyloading)
def chroma_query_loader():
    from .chroma_query import ChromaQuery
    return Injector().get(ChromaQuery)

@model_loader(lazy=blackboxConf.lazyloading)
def chroma_upsert_loader():
    from .chroma_upsert import ChromaUpsert
    return Injector().get(ChromaUpsert)

@model_loader(lazy=blackboxConf.lazyloading)
def chroma_chat_load():
    from .chroma_chat import ChromaChat
    return Injector().get(ChromaChat)

@singleton
class BlackboxFactory:
    models = {}

    @inject
    def __init__(self,) -> None:
        self.models["asr"] = asr_loader
        self.models["tts"] = tts_loader
        self.models["sentiment_engine"] = sentiment_loader
        self.models["emotion"] = emotion_loader
        self.models["fastchat"] = fastchat_loader
        self.models["audio_chat"] = audio_chat_loader
        self.models["g2e"] = g2e_loader
        self.models["text_and_image"] = text_and_image_loader
        self.models["chroma_query"] = chroma_query_loader
        self.models["chroma_upsert"] = chroma_upsert_loader
        self.models["chroma_chat"] = chroma_chat_load
        self.models["melotts"] = melotts_loader
        self.models["vlms"] = vlms_loader
        self.models["chat"] = chat_loader

    def __call__(self, *args, **kwargs):
        return self.processing(*args, **kwargs)

    def get_blackbox(self, blackbox_name: str) -> Blackbox:
        model = self.models.get(blackbox_name)
        if model is None:
            raise ValueError("Invalid Blackbox Type...")
        return model()