src/unlimited_ocr_rdna4/model_store.py

from __future__ import annotations

import hashlib
import json
import os
import shutil
import stat
import tempfile
from dataclasses import asdict, dataclass
from fcntl import LOCK_EX, flock
from pathlib import Path
from typing import BinaryIO

from .constants import (
    MODEL_DOWNLOAD_SHA256,
    MODEL_MANIFEST_FILE,
    MODEL_PAYLOAD_SHA256,
    MODEL_REPO,
    MODEL_REVISION,
    MODEL_UPSTREAM_BACKUP_FILE,
    MODEL_WEIGHT_BYTES,
    MODEL_WEIGHT_FILE,
    MODEL_WEIGHT_SHA256,
    PATCHED_MODEL_CODE_SHA256,
    UPSTREAM_MODEL_CODE_SHA256,
    VERIFIED_ROCM_VERSION,
    VERIFIED_TORCH_VERSION,
)
from .errors import ModelIntegrityError


@dataclass(frozen=True)
class ModelStatus:
    model_dir: str
    revision: str
    prepared: bool
    weight_bytes: int
    weight_sha256: str | None
    model_code_sha256: str
    reused: bool = False


def _regular_file(path: Path) -> bool:
    try:
        return stat.S_ISREG(path.lstat().st_mode)
    except FileNotFoundError:
        return False


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(4 * 1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def sha256_text(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def _replace_exact(text: str, old: str, new: str, expected: int) -> str:
    found = text.count(old)
    if found != expected:
        raise ModelIntegrityError(f"model patch context mismatch: expected {expected} occurrence(s), found {found}")
    return text.replace(old, new)


def patch_model_source_text(text: str) -> str:
    current_hash = sha256_text(text)
    if current_hash == PATCHED_MODEL_CODE_SHA256:
        return text
    if current_hash != UPSTREAM_MODEL_CODE_SHA256:
        raise ModelIntegrityError(
            "refusing to patch unknown model code; expected pinned Baidu revision "
            f"{MODEL_REVISION}, got SHA-256 {current_hash}"
        )

    text = _replace_exact(
        text,
        "from .modeling_deepseekv2 import DeepseekV2Model, DeepseekV2ForCausalLM",
        "import ast\n\nfrom .modeling_deepseekv2 import DeepseekV2Model, DeepseekV2ForCausalLM",
        1,
    )
    replacements = (
        ("cor_list = eval(ref_text[2])", "cor_list = ast.literal_eval(ref_text[2])"),
        ("lines = eval(outputs)['Line']['line']", "lines = ast.literal_eval(outputs)['Line']['line']"),
        (
            "line_type = eval(outputs)['Line']['line_type']",
            "line_type = ast.literal_eval(outputs)['Line']['line_type']",
        ),
        (
            "endpoints = eval(outputs)['Line']['line_endpoint']",
            "endpoints = ast.literal_eval(outputs)['Line']['line_endpoint']",
        ),
        ("p0 = eval(line.split(' -- ')[0])", "p0 = ast.literal_eval(line.split(' -- ')[0])"),
        ("p1 = eval(line.split(' -- ')[-1])", "p1 = ast.literal_eval(line.split(' -- ')[-1])"),
        ("(x, y) = eval(endpoint.split(': ')[1])", "(x, y) = ast.literal_eval(endpoint.split(': ')[1])"),
    )
    for old, new in replacements:
        text = _replace_exact(text, old, new, 1)

    text = _replace_exact(
        text,
        "images_seq_mask[idx].unsqueeze(-1).cuda()",
        "images_seq_mask[idx].unsqueeze(-1).to(inputs_embeds.device)",
        1,
    )
    text = _replace_exact(
        text,
        "                input_ids=input_ids.unsqueeze(0).cuda(),\n",
        "                input_ids=input_ids.unsqueeze(0).cuda(),\n"
        "                attention_mask=torch.ones_like(input_ids.unsqueeze(0), device='cuda'),\n",
        3,
    )
    text = _replace_exact(
        text,
        "                eos_token_id=tokenizer.eos_token_id,\n",
        "                eos_token_id=tokenizer.eos_token_id,\n                pad_token_id=tokenizer.eos_token_id,\n",
        3,
    )

    patched_hash = sha256_text(text)
    if PATCHED_MODEL_CODE_SHA256 != "__TO_BE_FILLED__" and patched_hash != PATCHED_MODEL_CODE_SHA256:
        raise ModelIntegrityError(f"patched model code hash mismatch: {patched_hash}")
    return text


def patch_model_file(model_dir: Path) -> str:
    source = model_dir / "modeling_unlimitedocr.py"
    text = source.read_text(encoding="utf-8")
    patched = patch_model_source_text(text)
    backup = model_dir / MODEL_UPSTREAM_BACKUP_FILE
    if not backup.exists() and sha256_text(text) == UPSTREAM_MODEL_CODE_SHA256:
        shutil.copy2(source, backup)
    temporary = source.with_suffix(".py.new")
    temporary.write_text(patched, encoding="utf-8")
    os.replace(temporary, source)
    return sha256_text(patched)


def _manifest_payload() -> dict[str, object]:
    return {
        "schema_version": 1,
        "source": MODEL_REPO,
        "revision": MODEL_REVISION,
        "files": dict(sorted(MODEL_PAYLOAD_SHA256.items())),
        "verified_rocm": VERIFIED_ROCM_VERSION,
        "verified_torch": VERIFIED_TORCH_VERSION,
    }


def _write_manifest(model_dir: Path) -> None:
    manifest = model_dir / MODEL_MANIFEST_FILE
    descriptor, temporary_name = tempfile.mkstemp(prefix=f".{MODEL_MANIFEST_FILE}.", dir=model_dir)
    try:
        with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
            json.dump(_manifest_payload(), handle, indent=2, sort_keys=True)
            handle.write("\n")
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(temporary_name, manifest)
    finally:
        Path(temporary_name).unlink(missing_ok=True)


def _verify_payload(model_dir: Path, expected_hashes: dict[str, str]) -> dict[str, str]:
    if model_dir.is_symlink() or not model_dir.is_dir():
        raise ModelIntegrityError(f"model path must be a real directory, not a symlink: {model_dir}")

    missing = [name for name in expected_hashes if not _regular_file(model_dir / name)]
    if missing:
        raise ModelIntegrityError(f"model is incomplete at {model_dir}; missing regular files: {', '.join(missing)}")

    weight = model_dir / MODEL_WEIGHT_FILE
    weight_bytes = weight.stat().st_size
    if weight_bytes != MODEL_WEIGHT_BYTES:
        raise ModelIntegrityError(f"weight size mismatch: expected {MODEL_WEIGHT_BYTES}, got {weight_bytes}")

    actual_hashes: dict[str, str] = {}
    for name, expected in expected_hashes.items():
        actual = sha256_file(model_dir / name)
        if actual != expected:
            raise ModelIntegrityError(f"model file SHA-256 mismatch for {name}: expected {expected}, got {actual}")
        actual_hashes[name] = actual
    return actual_hashes


def _verify_exact_layout(model_dir: Path) -> None:
    if model_dir.is_symlink() or not model_dir.is_dir():
        raise ModelIntegrityError(f"model path must be a real directory, not a symlink: {model_dir}")
    expected = set(MODEL_PAYLOAD_SHA256) | {MODEL_MANIFEST_FILE}
    actual = {entry.name for entry in model_dir.iterdir()}
    unexpected = sorted(actual - expected)
    missing = sorted(expected - actual)
    if unexpected or missing:
        details = []
        if missing:
            details.append(f"missing: {', '.join(missing)}")
        if unexpected:
            details.append(f"unexpected: {', '.join(unexpected)}")
        raise ModelIntegrityError(f"model layout mismatch at {model_dir}; {'; '.join(details)}")


def verify_model(model_dir: Path, *, full_weight_hash: bool = True) -> ModelStatus:
    del full_weight_hash  # Kept for API compatibility; security verification is always complete.
    _verify_exact_layout(model_dir)
    hashes = _verify_payload(model_dir, MODEL_PAYLOAD_SHA256)

    manifest_path = model_dir / MODEL_MANIFEST_FILE
    if not _regular_file(manifest_path):
        raise ModelIntegrityError(f"model manifest is missing or not a regular file: {manifest_path}")
    try:
        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        raise ModelIntegrityError(f"invalid model manifest at {manifest_path}: {exc}") from exc
    if manifest != _manifest_payload():
        raise ModelIntegrityError(f"model manifest does not match the pinned runtime contract: {manifest_path}")

    code_hash = hashes["modeling_unlimitedocr.py"]
    weight_hash = hashes[MODEL_WEIGHT_FILE]
    weight_bytes = (model_dir / MODEL_WEIGHT_FILE).stat().st_size

    return ModelStatus(
        model_dir=str(model_dir),
        revision=MODEL_REVISION,
        prepared=True,
        weight_bytes=weight_bytes,
        weight_sha256=weight_hash,
        model_code_sha256=code_hash,
    )


def _remove_snapshot_metadata(model_dir: Path) -> None:
    metadata = model_dir / ".cache"
    if metadata.is_symlink():
        raise ModelIntegrityError(f"refusing symlinked snapshot metadata: {metadata}")
    if metadata.exists():
        shutil.rmtree(metadata)


def _upgrade_existing_model(model_dir: Path) -> ModelStatus:
    allowed = set(MODEL_PAYLOAD_SHA256) | {MODEL_MANIFEST_FILE, ".cache"}
    actual = {entry.name for entry in model_dir.iterdir()}
    unexpected = sorted(actual - allowed)
    if unexpected:
        raise ModelIntegrityError(f"refusing unexpected entries in prepared model: {', '.join(unexpected)}")
    _verify_payload(model_dir, MODEL_PAYLOAD_SHA256)
    _remove_snapshot_metadata(model_dir)
    _write_manifest(model_dir)
    return verify_model(model_dir)


def _lock_handle(path: Path) -> BinaryIO:
    descriptor = os.open(path, os.O_RDWR | os.O_CREAT | os.O_CLOEXEC | os.O_NOFOLLOW, 0o600)
    handle = os.fdopen(descriptor, "a+b")
    flock(handle.fileno(), LOCK_EX)
    return handle


def prepare_model(model_dir: Path, cache_dir: Path, *, dry_run: bool = False) -> ModelStatus:
    model_dir = model_dir.expanduser().absolute()
    cache_dir = cache_dir.expanduser().resolve()
    if dry_run:
        return ModelStatus(
            model_dir=str(model_dir),
            revision=MODEL_REVISION,
            prepared=False,
            weight_bytes=MODEL_WEIGHT_BYTES,
            weight_sha256=MODEL_WEIGHT_SHA256,
            model_code_sha256=PATCHED_MODEL_CODE_SHA256,
        )

    if model_dir.is_symlink():
        raise ModelIntegrityError(f"model destination must not be a symlink: {model_dir}")
    model_dir.parent.mkdir(parents=True, exist_ok=True)
    cache_dir.mkdir(parents=True, exist_ok=True)

    try:
        from huggingface_hub import snapshot_download
    except ImportError as exc:
        raise ModelIntegrityError("huggingface-hub is missing; install the package dependencies") from exc

    lock_path = model_dir.parent / f".{model_dir.name}.prepare.lock"
    with _lock_handle(lock_path):
        if model_dir.exists():
            if not model_dir.is_dir():
                raise ModelIntegrityError(f"model destination exists and is not a directory: {model_dir}")
            status = _upgrade_existing_model(model_dir)
            return ModelStatus(**{**asdict(status), "reused": True})

        partial = Path(tempfile.mkdtemp(prefix=f".{model_dir.name}.partial-", dir=model_dir.parent))
        try:
            snapshot_download(
                repo_id=MODEL_REPO,
                revision=MODEL_REVISION,
                local_dir=partial,
                cache_dir=cache_dir / "huggingface",
                allow_patterns=sorted(MODEL_DOWNLOAD_SHA256),
            )
            _verify_payload(partial, MODEL_DOWNLOAD_SHA256)
            patch_model_file(partial)
            _remove_snapshot_metadata(partial)
            _write_manifest(partial)
            status = verify_model(partial)
            partial.rename(model_dir)
        except Exception:
            shutil.rmtree(partial, ignore_errors=True)
            raise
    return ModelStatus(**{**asdict(status), "model_dir": str(model_dir)})