from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

from PythonModels.components.amesim_pneumatic import (
    HELIUM_PNEUMATIC_GAS,
    AmesimPneumaticGas,
)
from PythonModels.components.amesim_pneumatic_line import AmesimPnl0001Pipe
from PythonModels.systems.test_mql_line_parameters import (
    TestMqlPnl0001Spec,
    load_test_mql_pnl0001_specs,
)


@dataclass(frozen=True)
class TestMqlPnl0001Assembly:
    specs: tuple[TestMqlPnl0001Spec, ...]
    lines: dict[str, AmesimPnl0001Pipe]

    def spec(self, alias: str) -> TestMqlPnl0001Spec:
        for spec in self.specs:
            if spec.alias == alias:
                return spec
        raise KeyError(alias)


def build_test_mql_pnl0001_assembly(
    archive_path: str | Path,
    *,
    gas: AmesimPneumaticGas = HELIUM_PNEUMATIC_GAS,
) -> TestMqlPnl0001Assembly:
    specs = load_test_mql_pnl0001_specs(archive_path)
    lines = {
        spec.alias: AmesimPnl0001Pipe(
            name=spec.alias,
            diameter_mm=spec.diameter_mm,
            length_m=spec.length_m,
            relative_roughness=spec.relative_roughness,
            polytropic_constant=spec.polytropic_constant,
            heat_transfer_coefficient=spec.heat_transfer_coefficient,
            external_temperature_k=spec.external_temperature_k,
            gas=gas,
            p0=spec.initial_absolute_pressure_pa,
            T0=spec.initial_temperature_k,
        )
        for spec in specs
    }
    return TestMqlPnl0001Assembly(specs=specs, lines=lines)