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,
    *,
    time_stop_s: float | None = None,
) -> AmesimResults:
    return load_amesim_results_from_archive(
        archive_path=archive_path,
        var_member=None,
        results_member=None,
        time_stop_s=time_stop_s,
    )


def load_amesim_results_from_archive(
    *,
    archive_path: str | Path,
    var_member: str | None,
    results_member: str | None,
    time_stop_s: float | None = None,
) -> AmesimResults:
    with tarfile.open(archive_path) as archive:
        var_member, results_member = _resolve_result_members(
            archive,
            var_member=var_member,
            results_member=results_member,
        )
        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()
        variables = tuple(
            _parse_variable_line(index, line) for index, line in enumerate(var_lines)
        )
        if time_stop_s is not None:
            return _parse_amesim_results_window(
                results_file,
                variables,
                time_stop_s=time_stop_s,
            )
        results_data = results_file.read()
    return parse_amesim_results_bytes(results_data, variables)


def _resolve_result_members(
    archive: tarfile.TarFile,
    *,
    var_member: str | None,
    results_member: str | None,
) -> tuple[str, str]:
    member_names = set(archive.getnames())
    if var_member is not None or results_member is not None:
        if var_member is None or results_member is None:
            raise AmesimResultsError(
                "var_member and results_member must either both be set or both be omitted."
            )
        return var_member, results_member

    preferred = ("test_mql_.var", "test_mql_.results")
    if preferred[0] in member_names and preferred[1] in member_names:
        return preferred

    pairs = sorted(
        (name, f"{name[:-4]}.results")
        for name in member_names
        if name.endswith(".var") and f"{name[:-4]}.results" in member_names
    )
    if len(pairs) != 1:
        raise AmesimResultsError(
            "Unable to identify a unique AMESim .var/.results member pair."
        )
    return pairs[0]


def _parse_amesim_results_window(
    results_file,
    variables: tuple[AmesimVariable, ...],
    *,
    time_stop_s: float,
) -> AmesimResults:
    header = results_file.read(8)
    if len(header) < 8:
        raise AmesimResultsError("AMESim results data is too small.")
    point_count, encoded_saved_variable_count = struct.unpack("<2i", header)
    saved_variable_count = abs(encoded_saved_variable_count)
    if point_count <= 0 or saved_variable_count <= 0:
        raise AmesimResultsError("Invalid AMESim results header.")

    mapping_data = results_file.read(saved_variable_count * 4)
    if len(mapping_data) != saved_variable_count * 4:
        raise AmesimResultsError("AMESim results variable mapping is truncated.")
    saved_variable_indices = struct.unpack(
        f"<{saved_variable_count}i",
        mapping_data,
    )
    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
    row_byte_count = row_length * 8
    times: list[float] = []
    series_lists: dict[str, list[float]] = {}
    saved_paths: list[tuple[int, str]] = []
    for column, variable_index in enumerate(saved_variable_indices, start=1):
        data_path = variables[variable_index].data_path
        if data_path is None:
            continue
        series_lists[data_path] = []
        saved_paths.append((column, data_path))

    for _row_index in range(point_count):
        row = results_file.read(row_byte_count)
        if len(row) != row_byte_count:
            raise AmesimResultsError("AMESim results matrix is truncated.")
        time_s = struct.unpack_from("<d", row, 0)[0]
        times.append(time_s)
        for column, data_path in saved_paths:
            series_lists[data_path].append(
                struct.unpack_from("<d", row, column * 8)[0]
            )
        # Keep one real sample after the requested stop so endpoint finite
        # differences do not silently fall back to a backward-only slope.
        if time_s > time_stop_s + 1.0e-12:
            break

    return AmesimResults(
        times=tuple(times),
        variables=variables,
        saved_variable_indices=tuple(saved_variable_indices),
        series_by_data_path={
            data_path: tuple(values) for data_path, values in series_lists.items()
        },
        final_values_by_data_path={},
    )


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,
    )