from typing import Dict, MutableMapping, cast

from qm.grpc.qm.pb import inc_qua_config_pb2
from qm.utils.protobuf_utils import which_one_of
from qm.api.models.server_details import ConnectionDetails
from qm.api.v2.job_api.element_api import ElementApi, JobElement
from qm.api.v2.job_api.element_input_api import InputConfigType, create_element_input_class
from qm.api.v2.job_api.element_port_api import MwOutputApi, AnalogOutputApi, DigitalInputApi


class JobElementNotFound(KeyError):
    def __init__(self, key: str):
        self._key = key

    def __str__(self) -> str:
        return f"Element with the key {self._key} was not found."


class JobElementsDB(Dict[str, JobElement]):
    def __missing__(self, key: str) -> None:
        raise JobElementNotFound(key)

    @classmethod
    def init_from_data(
        cls,
        elements: MutableMapping[str, inc_qua_config_pb2.QuaConfig.ElementDec],
        connection_details: ConnectionDetails,
        job_id: str,
    ) -> "JobElementsDB":
        res = {
            element_name: cls._create_single_element(element_name, element_config, connection_details, job_id)
            for element_name, element_config in elements.items()
        }
        return cls(res)

    @staticmethod
    def _create_single_element(
        element_name: str,
        element_config: inc_qua_config_pb2.QuaConfig.ElementDec,
        connection_details: ConnectionDetails,
        job_id: str,
    ) -> JobElement:
        input_config = cast(InputConfigType, which_one_of(element_config, "element_inputs_one_of")[1])
        input_api_class = create_element_input_class(input_config)

        mw_output_api = None
        outputs = which_one_of(element_config, "element_outputs_one_of")[1]
        if isinstance(outputs, inc_qua_config_pb2.QuaConfig.MultipleOutputs):
            output_names = list(outputs.port_references)
        elif isinstance(outputs, inc_qua_config_pb2.QuaConfig.MicrowaveOutputPortReference):
            output_names = list(element_config.outputs)
            mw_output_api = MwOutputApi(connection_details, job_id, element_name)
        else:
            output_names = list(element_config.outputs)

        outputs_apis = {
            output_name: AnalogOutputApi(connection_details, job_id, element_name, output_name)
            for output_name in output_names
        }
        digital_inputs_apis = {
            input_name: DigitalInputApi(connection_details, job_id, element_name, input_name)
            for input_name in element_config.digitalInputs
        }
        return JobElement(
            api=ElementApi(connection_details, job_id, element_name),
            input_api=input_api_class(connection_details, job_id, element_name),
            outputs_apis=outputs_apis,
            digital_inputs_apis=digital_inputs_apis,
            microwave_output_api=mw_output_api,
        )
