"""PlatformIO hooks for deterministic assets and partition size limits."""

from __future__ import annotations

import csv
import subprocess
from pathlib import Path

Import("env")


PROJECT_DIRECTORY = Path(env.subst("$PROJECT_DIR"))
BUILD_DIRECTORY = Path(env.subst("$BUILD_DIR"))
PARTITION_FILE = PROJECT_DIRECTORY / "partitions.csv"
ASSET_SCRIPT = PROJECT_DIRECTORY / "scripts" / "generate_web_assets.py"
FILESYSTEM_WRAPPER = PROJECT_DIRECTORY / "scripts" / "reproducible_mklittlefs.py"


def parse_size(value: str) -> int:
    value = value.strip()
    suffixes = {"K": 1024, "M": 1024 * 1024}
    if value[-1:].upper() in suffixes:
        return int(value[:-1], 0) * suffixes[value[-1].upper()]
    return int(value, 0)


def partition_sizes() -> dict[str, int]:
    sizes: dict[str, int] = {}
    with PARTITION_FILE.open(newline="", encoding="utf-8") as partition_stream:
        for row in csv.reader(partition_stream):
            if not row or row[0].lstrip().startswith("#"):
                continue
            if len(row) < 5:
                raise RuntimeError(f"Invalid partition row: {row}")
            sizes[row[0].strip()] = parse_size(row[4])
    return sizes


def generate_assets(source, target, env) -> None:
    command = [
        env.subst("$PYTHONEXE"),
        str(ASSET_SCRIPT),
        "--source",
        str(PROJECT_DIRECTORY / "web"),
        "--output",
        str(PROJECT_DIRECTORY / "data"),
    ]
    subprocess.run(command, check=True)


def enforce_size(path: Path, partition_name: str) -> None:
    limits = partition_sizes()
    if partition_name not in limits:
        raise RuntimeError(f"Missing required partition: {partition_name}")
    actual_size = path.stat().st_size
    maximum_size = limits[partition_name]
    if actual_size > maximum_size:
        raise RuntimeError(
            f"{path.name} is {actual_size} bytes and exceeds the "
            f"{partition_name} partition limit of {maximum_size} bytes"
        )
    print(
        f"maze32 size check: {path.name} uses {actual_size}/{maximum_size} bytes "
        f"in partition {partition_name}"
    )


def check_application(source, target, env) -> None:
    enforce_size(application_image, "factory")


def check_filesystem(source, target, env) -> None:
    enforce_size(filesystem_image, "littlefs")


filesystem_image = BUILD_DIRECTORY / "littlefs.bin"
application_image = BUILD_DIRECTORY / f"{env.subst('$PROGNAME')}.bin"
real_filesystem_tool = env.subst("$MKFSTOOL")

env["ENV"]["MAZE32_REAL_MKFSTOOL"] = real_filesystem_tool
env["ENV"]["MAZE32_BUILD_DIR"] = str(BUILD_DIRECTORY)
env.Replace(MKFSTOOL=str(FILESYSTEM_WRAPPER))

env.AddPreAction("$BUILD_DIR/littlefs.bin", generate_assets)
env.AddPostAction("$BUILD_DIR/littlefs.bin", check_filesystem)
env.AddPostAction("$BUILD_DIR/${PROGNAME}.bin", check_application)
env.AddCustomTarget(
    "webassets",
    None,
    generate_assets,
    title="Generate deterministic web assets",
    description="Compress web sources into the LittleFS data directory",
)