from __future__ import annotations

from collections import Counter
from dataclasses import dataclass

from PythonModels.reporting.amesim_results import AmesimResults
from PythonModels.reporting.test_mql_observations import (
    TestMqlObservationCatalog,
    build_test_mql_observation_catalog,
)
from PythonModels.reporting.test_mql_variables import TestMqlVariableBinding


@dataclass(frozen=True)
class TestMqlOutputSignal:
    data_path: str
    domain: str
    owner_alias: str
    owner_kind: str
    submodel: str
    signal_name: str
    units: str | None
    amesim_index: int
    saved: bool


@dataclass(frozen=True)
class TestMqlOutputSchema:
    signals: tuple[TestMqlOutputSignal, ...]

    @property
    def signal_count(self) -> int:
        return len(self.signals)

    def by_data_path(self, data_path: str) -> TestMqlOutputSignal:
        for signal in self.signals:
            if signal.data_path == data_path:
                return signal
        raise KeyError(data_path)

    def data_paths(self) -> tuple[str, ...]:
        return tuple(signal.data_path for signal in self.signals)

    def data_paths_by_domain(self, domain: str) -> tuple[str, ...]:
        return tuple(signal.data_path for signal in self.signals if signal.domain == domain)

    def counts_by_domain(self) -> dict[str, int]:
        return dict(Counter(signal.domain for signal in self.signals))

    def counts_by_submodel(self) -> dict[str, int]:
        return dict(Counter(signal.submodel for signal in self.signals))

    def counts_by_owner_kind(self) -> dict[str, int]:
        return dict(Counter(signal.owner_kind for signal in self.signals))

    def counts_by_units(self) -> dict[str | None, int]:
        return dict(Counter(signal.units for signal in self.signals))


def build_test_mql_output_schema(
    results: AmesimResults,
    *,
    observation_catalog: TestMqlObservationCatalog | None = None,
) -> TestMqlOutputSchema:
    observation_catalog = observation_catalog or build_test_mql_observation_catalog(results)
    domain_by_data_path = _domain_by_data_path(observation_catalog)
    signals = []
    for data_path in sorted(domain_by_data_path):
        variable = observation_catalog.variable_catalog.by_data_path(data_path)
        signals.append(_signal_from_variable(variable, domain_by_data_path[data_path]))
    return TestMqlOutputSchema(signals=tuple(signals))


def _domain_by_data_path(
    observation_catalog: TestMqlObservationCatalog,
) -> dict[str, str]:
    domain_by_data_path = {}
    for domain, data_paths in observation_catalog.data_paths_by_domain().items():
        for data_path in data_paths:
            if data_path in domain_by_data_path:
                raise ValueError(f"Data_Path {data_path!r} is assigned to multiple domains.")
            domain_by_data_path[data_path] = domain
    return domain_by_data_path


def _signal_from_variable(
    variable: TestMqlVariableBinding,
    domain: str,
) -> TestMqlOutputSignal:
    return TestMqlOutputSignal(
        data_path=variable.data_path,
        domain=domain,
        owner_alias=variable.owner_alias,
        owner_kind=variable.owner_kind,
        submodel=variable.submodel,
        signal_name=variable.signal_name,
        units=variable.units,
        amesim_index=variable.index,
        saved=variable.saved,
    )