# SPDX-FileCopyrightText: 2026 Tomasz Kramkowski <tomasz@kramkow.ski>
# SPDX-License-Identifier: GPL-3.0-or-later

import argparse
import hashlib
import sys
import tempfile
import tomllib
from base64 import b85encode
from importlib.resources import files
from io import BytesIO
from json import dumps
from pathlib import Path
from typing import Annotated, BinaryIO, Self
from zlib import compress

from pydantic import BaseModel, ConfigDict, Field, field_validator

from .elf import Hook, compile_hooks

NonnegativeInt = Annotated[int, Field(strict=True, ge=0)]


class Compilation(BaseModel):
    model_config = ConfigDict(extra="forbid")

    sources: list[Path] = Field(min_length=1)
    imports: dict[str, NonnegativeInt] = Field(default_factory=dict)
    hooks: list[Hook] = Field(min_length=1)


class FilePatch(BaseModel):
    model_config = ConfigDict(val_json_bytes="hex", extra="forbid")

    path: Path
    sha256: bytes = Field(min_length=32, max_length=32)
    edits: list[tuple[NonnegativeInt, bytes]]
    compiled: Compilation | None = None


class PatchSet(BaseModel):
    model_config = ConfigDict(extra="forbid")

    target: str
    targets: list[FilePatch]

    @field_validator("targets")
    @classmethod
    def unique_paths(cls, targets: list[FilePatch]) -> list[FilePatch]:
        paths = [target.path for target in targets]
        if len(set(paths)) != len(paths):
            raise ValueError("duplicate patch target")
        return targets

    @classmethod
    def parse(cls, config: dict, directory: Path) -> Self:
        result = cls.model_validate(config)
        for target in result.targets:
            if target.compiled is not None:
                target.compiled.sources = [
                    directory / source for source in target.compiled.sources
                ]
        return result


def compile_patch(config: FilePatch, source: BinaryIO, output: Path) -> dict:
    source.seek(0)
    data = source.read()
    if hashlib.sha256(data).digest() != config.sha256:
        raise ValueError(f"input SHA-256 mismatch: {config.path}")
    edits = list(config.edits)
    if config.compiled is not None:
        compiled = config.compiled
        edits += compile_hooks(
            source, compiled.sources, compiled.imports, compiled.hooks, output
        )
    end = 0
    for offset, replacement in sorted(edits):
        if offset < end:
            raise ValueError("overlapping edits")
        end = offset + len(replacement)
    patched = BytesIO(data)
    for offset, replacement in edits:
        patched.seek(offset)
        patched.write(replacement)
    return {
        "original_hash": config.sha256,
        "patched_hash": hashlib.sha256(patched.getbuffer()).digest(),
        "edits": edits,
    }


def render(patches: dict[str, dict]) -> str:
    template = files("steamvr_patches").joinpath("template.py").read_text()
    placeholder = "PATCH_DATA = {}"
    if template.count(placeholder) != 1:
        raise ValueError("missing or duplicate runtime patch-table placeholder")
    data = dumps(patches, default=bytes.hex, separators=(",", ":")).encode()
    encoded = b85encode(compress(data, level=9))
    return template.replace(placeholder, f"PATCH_DATA = decode_patch_data({encoded!r})")


def generate(patch_data: Path, input_root: Path) -> str:
    with (patch_data / "patches.toml").open("rb") as source:
        definition = PatchSet.parse(tomllib.load(source), patch_data)
    patches = {}
    with tempfile.TemporaryDirectory(prefix="generate-patcher-") as temporary:
        for index, config in enumerate(definition.targets):
            with (input_root / config.path).open("rb") as source:
                patches[config.path.as_posix()] = compile_patch(
                    config, source, Path(temporary) / str(index)
                )
    return render(patches)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "patch_data", type=Path, help="directory containing patches.toml etc"
    )
    parser.add_argument(
        "input_root", type=Path, help="directory containing the unpatched target files"
    )
    args = parser.parse_args()
    sys.stdout.write(generate(args.patch_data, args.input_root))