from __future__ import annotations

import re
from collections import Counter
from dataclasses import dataclass

from app.simulation.reporting.amesim_results import AmesimResults, AmesimVariable
from app.simulation.examples.test_mql.system import COMPONENT_SPECS, CONNECTION_SPECS


_UNIT_RE = re.compile(r"\[([^\]]+)\]\s*$")


@dataclass(frozen=True)
class TestMqlVariableBinding:
    index: int
    data_path: str
    signal_name: str
    owner_alias: str
    owner_kind: str
    submodel: str
    label: str
    units: str | None
    saved: bool


@dataclass(frozen=True)
class TestMqlVariableCatalog:
    variables: tuple[TestMqlVariableBinding, ...]

    @property
    def data_path_count(self) -> int:
        return len(self.variables)

    @property
    def saved_data_path_count(self) -> int:
        return sum(1 for variable in self.variables if variable.saved)

    def by_data_path(self, data_path: str) -> TestMqlVariableBinding:
        for variable in self.variables:
            if variable.data_path == data_path:
                return variable
        raise KeyError(data_path)

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

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

    def data_paths_for_owner(self, owner_alias: str) -> tuple[str, ...]:
        return tuple(
            variable.data_path
            for variable in self.variables
            if variable.owner_alias == owner_alias
        )


def build_test_mql_variable_catalog(amesim_results: AmesimResults) -> TestMqlVariableCatalog:
    owner_map = _build_owner_map()
    saved_indices = set(amesim_results.saved_variable_indices)
    bindings = []
    for variable in amesim_results.variables:
        if variable.data_path is None:
            continue
        signal_name, owner_alias = split_data_path(variable.data_path)
        owner_kind, submodel = owner_map[owner_alias]
        bindings.append(
            TestMqlVariableBinding(
                index=variable.index,
                data_path=variable.data_path,
                signal_name=signal_name,
                owner_alias=owner_alias,
                owner_kind=owner_kind,
                submodel=submodel,
                label=variable.label,
                units=_extract_units(variable),
                saved=variable.index in saved_indices,
            )
        )
    return TestMqlVariableCatalog(variables=tuple(bindings))


def split_data_path(data_path: str) -> tuple[str, str]:
    if "@" not in data_path:
        raise ValueError(f"AMESim Data_Path does not contain an owner alias: {data_path!r}")
    signal_name, owner_alias = data_path.rsplit("@", 1)
    if not signal_name or not owner_alias:
        raise ValueError(f"Invalid AMESim Data_Path: {data_path!r}")
    return signal_name, owner_alias


def _build_owner_map() -> dict[str, tuple[str, str]]:
    owner_map = {
        str(spec["alias"]): ("component", str(spec["submodel"]))
        for spec in COMPONENT_SPECS
    }
    owner_map.update(
        {
            str(spec["alias"]): ("connection", str(spec["submodel"]))
            for spec in CONNECTION_SPECS
        }
    )
    return owner_map


def _extract_units(variable: AmesimVariable) -> str | None:
    match = _UNIT_RE.search(variable.label)
    if match is None:
        return None
    return match.group(1)