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,
    AmesimPnl0003Pipe,
    AmesimPnl00rPipe,
)
from PythonModels.systems.test_mql_line_parameters import (
    TestMqlPnl0001Spec,
    TestMqlPnl0003Spec,
    TestMqlPnl00rSpec,
    load_test_mql_pnl0001_specs,
    load_test_mql_pnl0003_specs,
    load_test_mql_pnl00r_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)


@dataclass(frozen=True)
class TestMqlPnl0003Assembly:
    specs: tuple[TestMqlPnl0003Spec, ...]
    lines: dict[str, AmesimPnl0003Pipe]

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


@dataclass(frozen=True)
class TestMqlPnl00rAssembly:
    specs: tuple[TestMqlPnl00rSpec, ...]
    lines: dict[str, AmesimPnl00rPipe]

    def spec(self, alias: str) -> TestMqlPnl00rSpec:
        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)


def build_test_mql_pnl0003_assembly(
    archive_path: str | Path,
    *,
    gas: AmesimPneumaticGas = HELIUM_PNEUMATIC_GAS,
) -> TestMqlPnl0003Assembly:
    specs = load_test_mql_pnl0003_specs(archive_path)
    lines = {
        spec.alias: AmesimPnl0003Pipe(
            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,
            p1_0=spec.initial_absolute_pressure_1_pa,
            T1_0=spec.initial_temperature_1_k,
            p2_0=spec.initial_absolute_pressure_2_pa,
            T2_0=spec.initial_temperature_2_k,
        )
        for spec in specs
    }
    return TestMqlPnl0003Assembly(specs=specs, lines=lines)


def build_test_mql_pnl00r_assembly(
    archive_path: str | Path,
    *,
    gas: AmesimPneumaticGas = HELIUM_PNEUMATIC_GAS,
) -> TestMqlPnl00rAssembly:
    specs = load_test_mql_pnl00r_specs(archive_path)
    lines = {
        spec.alias: AmesimPnl00rPipe(
            name=spec.alias,
            diameter_mm=spec.diameter_mm,
            length_m=spec.length_m,
            relative_roughness=spec.relative_roughness,
            gas=gas,
        )
        for spec in specs
    }
    return TestMqlPnl00rAssembly(specs=specs, lines=lines)