from __future__ import annotations

import json
import zlib
from collections.abc import Mapping
from pathlib import Path
from typing import Any

import msgpack

from bedit_gui.simulation_models import SimulationRoot

BES_MAGIC = b"BES\x00"
FILE_FORMAT_VERSION = 1
_VERSION_SIZE = 4


def load(path: str | Path) -> SimulationRoot:
    file_path = Path(path)
    data = _load_json(file_path) if file_path.suffix.lower() == ".json" else _load_bes(file_path)
    return SimulationRoot.from_data(data)


def save(root: SimulationRoot, path: str | Path) -> None:
    file_path = Path(path)
    if file_path.suffix.lower() == ".json":
        _save_json(root.to_data(), file_path)
    elif file_path.suffix.lower() == ".bes":
        _save_bes(root.to_data(), file_path)
    else:
        raise ValueError(f"unsupported simulation file extension {file_path.suffix!r}; expected '.json' or '.bes'")


def is_simulation_json(path: str | Path) -> bool:
    file_path = Path(path)
    if file_path.suffix.lower() != ".json":
        return False
    try:
        return _load_json(file_path).get("root_type") == "simulation_root"
    except (TypeError, ValueError):
        return False


def _load_json(path: Path) -> Mapping[str, Any]:
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, UnicodeError, json.JSONDecodeError) as exc:
        raise ValueError(f"could not read simulation JSON {path}: {exc}") from exc
    if not isinstance(data, Mapping):
        raise TypeError(f"simulation JSON {path} must contain an object")
    return data


def _save_json(data: Mapping[str, Any], path: Path) -> None:
    path.write_text(json.dumps({"file_format_version": FILE_FORMAT_VERSION, **data}, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def _load_bes(path: Path) -> Mapping[str, Any]:
    try:
        payload = path.read_bytes()
        header_end = len(BES_MAGIC) + _VERSION_SIZE
        if not payload.startswith(BES_MAGIC):
            raise ValueError("missing BES file header")
        if len(payload) < header_end:
            raise ValueError("truncated BES file header")
        version = int.from_bytes(payload[len(BES_MAGIC):header_end], "big")
        if version != FILE_FORMAT_VERSION:
            raise ValueError(f"unsupported BES file version {version}")
        data = msgpack.unpackb(zlib.decompress(payload[header_end:]), raw=False, strict_map_key=False)
    except (OSError, ValueError, zlib.error, msgpack.exceptions.MsgpackException) as exc:
        raise ValueError(f"could not read BES simulation {path}: {exc}") from exc
    if not isinstance(data, Mapping):
        raise TypeError(f"BES simulation {path} must contain a map")
    return data


def _save_bes(data: Mapping[str, Any], path: Path) -> None:
    encoded = zlib.compress(msgpack.packb(dict(data), use_bin_type=True))
    path.write_bytes(BES_MAGIC + FILE_FORMAT_VERSION.to_bytes(_VERSION_SIZE, "big") + encoded)