from __future__ import annotations

from typing import Union

from pydantic import BaseModel, Field, model_validator

from ..lang.datamodel_enum import (
    EnvironmentEnum,
    MetricClassEnum,
    MetricNameEnum,
    ServiceAPIEnum,
)


class ClientInfo(BaseModel):
    id: str = Field(
        default="",
        description="The id of the QoAClient, should be not be set manually, auto-generated by QoAClient",
    )
    name: str = Field(
        default="",
        description="The name of the QoAClient, for example, qoa_client1",
    )

    user_id: str = Field(
        default="",
        description="Can be set if you have external user_id, otherwise auto generated",
    )
    username: str = Field(default="", description="The name of the qoa_client user")
    instance_id: str = Field(
        default="",
        description="An UUID4 that is used to identify the instance, can be set manually in the config or through environment variable INSTANCE_ID, if not, auto-generated by QoAClient",
    )
    instance_name: str = Field(default="", description="The name of the instance")
    stage_id: str = ""
    functionality: str = ""
    application_name: str = ""
    role: str = ""
    run_id: str = ""
    environment: EnvironmentEnum = Field(
        default=EnvironmentEnum.edge,
        description="The environment where the qoa_client is run",
    )
    custom_info: dict | str = Field(
        default="", description="Additional information that you can add"
    )


class AMQPCollectorConfig(BaseModel):
    end_point: str
    exchange_name: str
    exchange_type: str
    in_routing_key: str
    in_queue: str


class AMQPConnectorConfig(BaseModel):
    end_point: str
    exchange_name: str
    exchange_type: str
    out_routing_key: str


class MQTTConnectorConfig(BaseModel):
    in_queue: str
    out_queue: str
    broker_url: str
    broker_port: int
    broker_keepalive: int
    client_id: str


class SocketConnectorConfig(BaseModel):
    host: str
    port: int


class SocketCollectorConfig(BaseModel):
    host: str
    port: int
    backlog: int
    bufsize: int


class PrometheusConnectorConfig(BaseModel):
    pass


class KafkaConnectorConfig(BaseModel):
    topic: str
    broker_url: str


class KafkaCollectorConfig(BaseModel):
    topic: str
    broker_url: str
    group_id: str
    auto_offset_reset: str = "earliest"
    poll_inteval: float = 1.0


class DebugConnectorConfig(BaseModel):
    silence: bool


# TODO: test if loading the config, the type of the config can be found
CollectorConfigClass = Union[
    AMQPCollectorConfig, SocketCollectorConfig, KafkaCollectorConfig, dict
]
ConnectorConfigClass = Union[
    AMQPConnectorConfig,
    SocketConnectorConfig,
    KafkaConnectorConfig,
    DebugConnectorConfig,
    dict,
]


class CollectorConfig(BaseModel):
    name: str
    collector_class: ServiceAPIEnum
    config: CollectorConfigClass


class ConnectorConfig(BaseModel):
    name: str
    connector_class: ServiceAPIEnum
    config: ConnectorConfigClass | None = None


class ClientConfig(BaseModel):
    client: ClientInfo
    registration_url: str | None = None
    collector: list[CollectorConfig] | None = None
    connector: list[ConnectorConfig] | None = None
    probes: list[ProbeConfig] | None = None

    @model_validator(mode="before")
    @classmethod
    def validate_probe_type(cls, values):
        probe_type_map: dict[str, type[ProbeConfig]] = {
            "process": ProcessProbeConfig,
            "docker": DockerProbeConfig,
            "system": SystemProbeConfig,
        }
        probes = values.get("probes", [])
        if probes:
            for i, probe_data in enumerate(probes):
                probe_type = probe_data.get("probe_type")
                if probe_type in probe_type_map:
                    probes[i] = probe_type_map[probe_type](**probe_data)
            values["probes"] = probes
        return values


class MetricConfig(BaseModel):
    name: MetricNameEnum
    metric_class: MetricClassEnum
    description: str | None = None
    default_value: int
    category: int


class GroupMetricConfig(BaseModel):
    name: str
    metric_configs: list[MetricConfig]


class ProbeConfig(BaseModel):
    probe_type: str
    frequency: int
    require_register: bool
    obs_service_url: str | None = None
    log_latency_flag: bool
    latency_logging_path: str | None = None
    environment: EnvironmentEnum = Field(
        default=EnvironmentEnum.edge, description="The environment where the probe run"
    )


class ProcessProbeConfig(ProbeConfig):
    probe_type: str = Field("process")
    pid: int | None = None


class SystemProbeConfig(ProbeConfig):
    probe_type: str = Field("system")
    node_name: str | None = None


class DockerProbeConfig(ProbeConfig):
    probe_type: str = Field("docker")
    container_list: list[str] = []


class JetsonSystemProbeConfig(ProbeConfig):
    probe_type: str = Field("jetson_sys")
    node_name: str | None


class JetsonProcessesProbeConfig(ProbeConfig):
    probe_type: str = Field("jetson_sys")
    node_name: str | None


class NodeAggregatorConfig(BaseModel):
    socket_collector_config: SocketCollectorConfig
    environment: EnvironmentEnum
    query_method: str
    data_separator: str
    unit_conversion: dict
    database_path: str


class ExporterConfig(BaseModel):
    host: str
    port: int
    node_aggregator: NodeAggregatorConfig


class OdopObsConfig(BaseModel):
    process: ProcessProbeConfig
    system: SystemProbeConfig
    probe_connector: SocketConnectorConfig
    exporter: ExporterConfig
