# Creates a directory in which to look up available agents

from typing import List
from src.simuleval_transcoder import SimulevalTranscoder
import json
import logging

logger = logging.getLogger("socketio_server_pubsub")

# fmt: off
M4T_P0_LANGS = [
    "eng",
    "arb", "ben", "cmn", "deu", "fil",
    "fra", "hin", "ind", "ita", "jpn",
    "kor", "nld", "por", "rus", "spa",
    "swh", "tha", "tur", "urd", "vie",
]
# fmt: on


class NoAvailableAgentException(Exception):
    pass


class AgentWithInfo:
    def __init__(
        self,
        agent,
        name: str,
        modalities: List[str],
        source_langs: List[str],
        target_langs: List[str],
        # Supported dynamic params are defined in StreamingTypes.ts
        dynamic_params: List[str] = [],
        description="",
    ):
        self.agent = agent
        self.name = name
        self.description = description
        self.modalities = modalities
        self.source_langs = source_langs
        self.target_langs = target_langs
        self.dynamic_params = dynamic_params

    def get_capabilities_for_json(self):
        return {
            "name": self.name,
            "description": self.description,
            "modalities": self.modalities,
            "sourceLangs": self.source_langs,
            "targetLangs": self.target_langs,
            "dynamicParams": self.dynamic_params,
        }

    @classmethod
    def load_from_json(cls, config: str):
        """
        Takes in JSON array of models to load in, e.g.
        [{"name": "s2s_m4t_emma-unity2_multidomain_v0.1", "description": "M4T model that supports simultaneous S2S and S2T", "modalities": ["s2t", "s2s"], "sourceLangs": ["arb", "ben", "hin", "ind", "ita", "jpn", "por", "rus", "spa", "swh", "tha", "tur", "urd", "vie"], "targetLangs": ["en"]},
        {"name": "s2s_m4t_expr-emma_v0.1", "description": "ES-EN expressive model that supports S2S and S2T", "modalities": ["s2t", "s2s"], "sourceLangs": ["arb", "ben", "hin", "ind", "ita", "jpn", "por", "rus", "spa", "swh", "tha", "tur", "urd", "vie"], "targetLangs": ["en"]}]
        """
        configs = json.loads(config)
        agents = []
        for config in configs:
            agent = SimulevalTranscoder.build_agent(config["name"])
            agents.append(
                AgentWithInfo(
                    agent=agent,
                    name=config["name"],
                    modalities=config["modalities"],
                    source_langs=config["sourceLangs"],
                    target_langs=config["targetLangs"],
                )
            )
        return agents


class SimulevalAgentDirectory:
    # Available models. These are the directories where the models can be found, and also serve as an ID for the model.
    # s2t:
    s2t_es_en_agent = "s2t_es-en_tt-waitk_multidomain"
    s2t_en_es_agent = "s2t_en-es_tt-waitk_multidomain"
    s2t_es_en_emma_agent = "s2t_es-en_emma_multidomain_v0.3"
    s2t_en_es_emma_agent = "s2t_en-es_emma_multidomain_v0.3"
    # s2s:
    s2s_es_en_agent = "s2s_es-en_tt-waitk-unity2_multidomain"
    s2s_es_en_emma_agent = "s2s_es-en_emma-unity2_multidomain_v0.2"
    s2s_m4t_expr_emma_agent = "s2s_m4t_expr-emma_v0.3"
    s2s_m4t_emma_agent = "s2s_m4t_emma-unity2_multidomain_v0.4"

    def __init__(self):
        self.agents = []
        self.did_build_and_add_agents = False

    def add_agent(self, agent: AgentWithInfo):
        self.agents.append(agent)

    def build_agent_if_available(self, model_id, config_name=None):
        agent = None
        try:
            if config_name is not None:
                agent = SimulevalTranscoder.build_agent(
                    model_id,
                    config_name=config_name,
                )
            else:
                agent = SimulevalTranscoder.build_agent(
                    model_id,
                )
        except Exception as e:
            logger.warning("Failed to build agent %s: %s" % (model_id, e))
            raise e

        return agent

    def build_and_add_agents(self, models_override=None):
        if self.did_build_and_add_agents:
            return

        if models_override is not None:
            agent_infos = AgentWithInfo.load_from_json(models_override)
            for agent_info in agent_infos:
                self.add_agent(agent_info)
        else:
            s2s_m4t_expr_agent = self.build_agent_if_available(
                SimulevalAgentDirectory.s2s_m4t_expr_emma_agent,
                config_name="vad_s2st_sc_24khz_main.yaml",
            )

            if s2s_m4t_expr_agent:
                self.add_agent(
                    AgentWithInfo(
                        agent=s2s_m4t_expr_agent,
                        name=SimulevalAgentDirectory.s2s_m4t_expr_emma_agent,
                        modalities=["s2t", "s2s"],
                        source_langs=M4T_P0_LANGS,
                        target_langs=["eng", "spa", "fra", "deu", "ita", "cmn"],
                        dynamic_params=["expressive"],
                        description="multilingual expressive model that supports S2S and S2T",
                    )
                )

        if len(self.agents) == 0:
            logger.error(
                "No agents were loaded. This likely means you are missing the actual model files specified in simuleval_agent_directory."
            )

        self.did_build_and_add_agents = True

    def get_agent(self, name):
        for agent in self.agents:
            if agent.name == name:
                return agent.agent
        return None

    def get_agent_or_throw(self, name):
        agent = self.get_agent(name)
        if agent is None:
            raise NoAvailableAgentException("No agent found with name= %s" % (name))
        return agent

    def get_agents_capabilities_list_for_json(self):
        return [agent.get_capabilities_for_json() for agent in self.agents]