from __future__ import annotations

import re
import struct
import tarfile
from dataclasses import dataclass
from pathlib import Path


class AmesimResultsError(ValueError):
    """Raised when AMESim result files cannot be parsed consistently."""


@dataclass(frozen=True)
class AmesimVariable:
    index: int
    label: str
    data_path: str | None
    param_id: int | None
    hidden: bool


@dataclass(frozen=True)
class AmesimResults:
    times: tuple[float, ...]
    variables: tuple[AmesimVariable, ...]
    saved_variable_indices: tuple[int, ...]
    series_by_data_path: dict[str, tuple[float, ...]]
    final_values_by_data_path: dict[str, float]

    @property
    def point_count(self) -> int:
        return len(self.times)

    @property
    def saved_variable_count(self) -> int:
        return len(self.saved_variable_indices)

    def series(self, data_path: str) -> tuple[float, ...]:
        return self.series_by_data_path[data_path]

    def final_value(self, data_path: str) -> float:
        return self.final_values_by_data_path[data_path]


_DATA_PATH_RE = re.compile(r"Data_Path=(\S+)")
_PARAM_ID_RE = re.compile(r"Param_Id=(\d+)")


def load_test_mql_amesim_results(archive_path: str | Path) -> AmesimResults:
    return load_amesim_results_from_archive(
        archive_path=archive_path,
        var_member="test_mql_.var",
        results_member="test_mql_.results",
    )


def load_amesim_results_from_archive(
    *,
    archive_path: str | Path,
    var_member: str,
    results_member: str,
) -> AmesimResults:
    with tarfile.open(archive_path) as archive:
        var_file = archive.extractfile(var_member)
        results_file = archive.extractfile(results_member)
        if var_file is None:
            raise AmesimResultsError(f"Missing AMESim variable member: {var_member}")
        if results_file is None:
            raise AmesimResultsError(f"Missing AMESim results member: {results_member}")
        var_lines = var_file.read().decode("latin1").splitlines()
        results_data = results_file.read()
    variables = tuple(_parse_variable_line(index, line) for index, line in enumerate(var_lines))
    return parse_amesim_results_bytes(results_data, variables)


def parse_amesim_results_bytes(
    results_data: bytes,
    variables: tuple[AmesimVariable, ...],
) -> AmesimResults:
    if len(results_data) < 8:
        raise AmesimResultsError("AMESim results data is too small.")
    point_count, encoded_saved_variable_count = struct.unpack_from("<2i", results_data, 0)
    saved_variable_count = abs(encoded_saved_variable_count)
    if point_count <= 0 or saved_variable_count <= 0:
        raise AmesimResultsError("Invalid AMESim results header.")

    mapping_offset = 8
    mapping_size = saved_variable_count * 4
    data_offset = mapping_offset + mapping_size
    saved_variable_indices = struct.unpack_from(
        f"<{saved_variable_count}i",
        results_data,
        mapping_offset,
    )
    if any(index < 0 or index >= len(variables) for index in saved_variable_indices):
        raise AmesimResultsError(
            "AMESim results variable mapping references unknown .var rows."
        )

    row_length = 1 + saved_variable_count
    main_value_count = point_count * row_length
    main_byte_count = main_value_count * 8
    main_end = data_offset + main_byte_count
    if main_end > len(results_data):
        raise AmesimResultsError("AMESim results matrix is truncated.")

    main_values = struct.unpack_from(f"<{main_value_count}d", results_data, data_offset)
    times = tuple(main_values[row * row_length] for row in range(point_count))
    series_by_data_path: dict[str, tuple[float, ...]] = {}
    for column, variable_index in enumerate(saved_variable_indices, start=1):
        variable = variables[variable_index]
        if variable.data_path is None:
            continue
        series_by_data_path[variable.data_path] = tuple(
            main_values[row * row_length + column]
            for row in range(point_count)
        )

    final_values_by_data_path: dict[str, float] = {}
    trailing_bytes = len(results_data) - main_end
    expected_final_bytes = (1 + len(variables)) * 8
    if trailing_bytes >= expected_final_bytes:
        final_values = struct.unpack_from(f"<{1 + len(variables)}d", results_data, main_end)
        for variable, value in zip(variables, final_values[1:]):
            if variable.data_path is not None:
                final_values_by_data_path[variable.data_path] = value

    return AmesimResults(
        times=times,
        variables=variables,
        saved_variable_indices=tuple(saved_variable_indices),
        series_by_data_path=series_by_data_path,
        final_values_by_data_path=final_values_by_data_path,
    )


def _parse_variable_line(index: int, line: str) -> AmesimVariable:
    data_path_match = _DATA_PATH_RE.search(line)
    param_id_match = _PARAM_ID_RE.search(line)
    label = line
    if data_path_match is not None:
        label = line[: data_path_match.start()].strip()
    return AmesimVariable(
        index=index,
        label=label,
        data_path=data_path_match.group(1) if data_path_match else None,
        param_id=int(param_id_match.group(1)) if param_id_match else None,
        hidden="HIDDEN" in line,
    )