"""Prepare immutable model/tokenizer assets during build, never during startup."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
from pathlib import Path

DEFAULT_MODEL = "intfloat/multilingual-e5-small"
DEFAULT_REVISION = "614241f622f53c4eeff9890bdc4f31cfecc418b3"
DEFAULT_ROOT = Path(__file__).resolve().parents[2] / "runtime_assets"


def asset_root() -> Path:
    return Path(os.getenv("ZP_RUNTIME_ASSETS", str(DEFAULT_ROOT)))


def manifest(root: Path | None = None) -> dict | None:
    path = (root or asset_root()) / "manifest.json"
    if not path.exists():
        if os.getenv("ZP_ASSETS_REQUIRED") == "1":
            raise RuntimeError(
                "Chýbajú runtime assets; zostavte image alebo spustite python -m tools.runtime.runtime_assets prepare."
            )
        return None
    return json.loads(path.read_text())


def configure_tokenizer() -> None:
    data = manifest()
    if data is not None:
        # Bundled immutable files take precedence over an empty user cache.
        os.environ["TIKTOKEN_CACHE_DIR"] = str(asset_root() / "tiktoken")


def local_model(model_name: str, revision: str | None) -> str | None:
    data = manifest()
    if data and data["embedding"]["model"] == model_name:
        if revision and revision != data["embedding"]["revision"]:
            raise RuntimeError(
                "Požadovaná revízia modelu nie je v image; vytvorte nový image s touto revíziou."
            )
        return str(asset_root() / "model")
    if data and os.getenv("ZP_ASSETS_REQUIRED") == "1":
        raise RuntimeError("Požadovaný embeddingový model nie je v image; vytvorte nový image.")
    return None


def verify(root: Path) -> dict:
    data = json.loads((root / "manifest.json").read_text())
    for relative, digest in data["files"].items():
        path = root / relative
        if not path.resolve().is_relative_to(root.resolve()):
            raise RuntimeError("Manifest obsahuje cestu mimo runtime assets")
        if not path.is_file() or hashlib.sha256(path.read_bytes()).hexdigest() != digest:
            raise RuntimeError(f"Chýbajúci alebo poškodený runtime asset: {relative}")
    if not any(name.startswith("tiktoken/") for name in data["files"]):
        raise RuntimeError("Manifest neobsahuje tokenizér")
    required = {"model/config.json", "model/modules.json", "model/tokenizer.json"}
    if not required.issubset(data["files"]) or not any(
        name.endswith(".safetensors") for name in data["files"]
    ):
        raise RuntimeError("Manifest neobsahuje kompletný embeddingový model")
    return data


def offline_probe(root: Path) -> None:
    """Fresh interpreter, with network calls forbidden even on an online builder."""
    environment = {
        **os.environ,
        "ZP_RUNTIME_ASSETS": str(root.resolve()),
        "ZP_ASSETS_REQUIRED": "1",
        "HF_HUB_OFFLINE": "1",
        "TRANSFORMERS_OFFLINE": "1",
        "HF_HOME": str(root / "unused-user-cache"),
        "TIKTOKEN_CACHE_DIR": str(root / "unused-token-cache"),
    }
    code = """
import socket
def denied(*args, **kwargs): raise RuntimeError('Network access during offline asset probe')
socket.socket.connect = denied
socket.socket.connect_ex = denied
socket.getaddrinfo = denied
from tools.runtime.runtime_assets import configure_tokenizer, manifest, local_model
configure_tokenizer()
import tiktoken
for encoding in manifest()['encodings']:
    assert tiktoken.get_encoding(encoding).encode('Overenie štartu bez siete.')
from sentence_transformers import SentenceTransformer
data = manifest()['embedding']
model = SentenceTransformer(local_model(data['model'], data['revision']), local_files_only=True)
assert model.encode(['query: overenie bez siete'], show_progress_bar=False).shape[0] == 1
print('Offline model/tokenizer probe: OK')
"""
    subprocess.run([sys.executable, "-c", code], env=environment, check=True)


def prepare(root: Path, model: str, revision: str | None, encodings: list[str]) -> dict:
    revision = revision or (DEFAULT_REVISION if model == DEFAULT_MODEL else "")
    if not re.fullmatch("[0-9a-f]{40}", revision):
        raise ValueError("Príprava modelu vyžaduje presnú 40-znakovú Git revíziu")
    if (root / "manifest.json").is_file():
        data = verify(root)
        if (
            data["embedding"] == {"model": model, "revision": revision}
            and data["encodings"] == encodings
        ):
            offline_probe(root)
            return data
        raise RuntimeError(
            "Assets patria inej konfigurácii; použite nový/prázdny pracovný adresár."
        )
    from huggingface_hub import snapshot_download
    from huggingface_hub.errors import LocalEntryNotFoundError

    try:
        source = snapshot_download(model, revision=revision, local_files_only=True)
    except (FileNotFoundError, LocalEntryNotFoundError):
        source = snapshot_download(
            model,
            revision=revision,
            allow_patterns=["*.json", "*.safetensors", "*.model", "*.txt"],
            ignore_patterns=["onnx/*", "openvino/*"],
        )
    root.mkdir(parents=True, exist_ok=True)
    shutil.copytree(
        source,
        root / "model",
        dirs_exist_ok=True,
        symlinks=False,
        ignore=shutil.ignore_patterns(
            "onnx", "openvino", "*.bin", "*.md", ".cache", ".gitattributes"
        ),
    )
    environment = {**os.environ, "TIKTOKEN_CACHE_DIR": str(root.resolve() / "tiktoken")}
    subprocess.run(
        [
            sys.executable,
            "-c",
            "import sys,tiktoken; [tiktoken.get_encoding(name) for name in sys.argv[1:]]",
            *encodings,
        ],
        env=environment,
        check=True,
    )
    files = {
        str(p.relative_to(root)): hashlib.sha256(p.read_bytes()).hexdigest()
        for folder in ["model", "tiktoken"]
        for p in sorted((root / folder).rglob("*"))
        if p.is_file()
    }
    data = {
        "schema_version": 1,
        "embedding": {"model": model, "revision": revision},
        "encodings": encodings,
        "files": files,
    }
    (root / "manifest.json").write_text(json.dumps(data, sort_keys=True, indent=2) + "\n")
    verify(root)
    offline_probe(root)
    return data


RERANKER_MODEL = "cross-encoder/mmarco-mMiniLMv2-L12-H384-v1"
RERANKER_REVISION = "1427fd652930e4ba29e8149678df786c240d8825"
STANZA_LANGUAGE = "sk"
STANZA_PROCESSORS = "tokenize,pos,lemma"


def _hashes(root: Path, folders) -> dict:
    return {
        str(p.relative_to(root)): hashlib.sha256(p.read_bytes()).hexdigest()
        for folder in folders
        for p in sorted((root / folder).rglob("*"))
        if p.is_file()
    }


def verify_extras(root: Path) -> dict:
    """Reranker and Stanza Slovak files unchanged since preparation."""
    path = root / "extras.json"
    if not path.is_file():
        raise RuntimeError("Chýba manifest doplnkových modelov (reranker, Stanza).")
    data = json.loads(path.read_text())
    if _hashes(root, ["reranker", "stanza"]) != data["files"]:
        raise RuntimeError("Doplnkové modely sa líšia od manifestu.")
    return data


def extras_probe(root: Path) -> None:
    """Load both models without network access."""
    code = f"""
import socket
def denied(*args, **kwargs): raise RuntimeError('Network access during offline asset probe')
socket.create_connection = denied
from sentence_transformers import CrossEncoder
model = CrossEncoder({str(root / 'reranker')!r}, local_files_only=True, device='cpu')
assert len(model.predict([('otázka', 'text')])) == 1
import stanza
nlp = stanza.Pipeline({STANZA_LANGUAGE!r}, dir={str(root / 'stanza')!r}, processors={STANZA_PROCESSORS!r},
                      download_method=None, verbose=False, use_gpu=False)
assert nlp('chatbotom').sentences[0].words[0].lemma == 'chatbot'
print('Offline reranker/Stanza probe: OK')
"""
    environment = {**os.environ, "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1"}
    subprocess.run([sys.executable, "-c", code], env=environment, check=True)


def prepare_extras(root: Path) -> dict:
    """Pinned cross-encoder and Stanza Slovak lemmatizer, copied into root."""
    if (root / "extras.json").is_file():
        data = verify_extras(root)
        if data["reranker"] == {"model": RERANKER_MODEL, "revision": RERANKER_REVISION}:
            extras_probe(root)
            return data
        raise RuntimeError("Doplnkové modely patria inej konfigurácii.")
    from huggingface_hub import snapshot_download
    from huggingface_hub.errors import LocalEntryNotFoundError
    try:
        source = snapshot_download(RERANKER_MODEL, revision=RERANKER_REVISION, local_files_only=True)
    except (FileNotFoundError, LocalEntryNotFoundError):
        source = snapshot_download(RERANKER_MODEL, revision=RERANKER_REVISION,
                                   allow_patterns=["*.json", "*.safetensors", "*.model", "*.txt"])
    root.mkdir(parents=True, exist_ok=True)
    shutil.copytree(source, root / "reranker", dirs_exist_ok=True, symlinks=False,
                    ignore=shutil.ignore_patterns("onnx", "openvino", "*.bin", "*.md", ".cache", ".gitattributes"))
    import stanza
    from stanza.resources.common import DEFAULT_MODEL_DIR
    cached = Path(DEFAULT_MODEL_DIR)
    if (cached / STANZA_LANGUAGE).is_dir() and (cached / "resources.json").is_file():
        (root / "stanza").mkdir(parents=True, exist_ok=True)
        shutil.copy2(cached / "resources.json", root / "stanza" / "resources.json")
        shutil.copytree(cached / STANZA_LANGUAGE, root / "stanza" / STANZA_LANGUAGE, dirs_exist_ok=True)
    else:
        stanza.download(STANZA_LANGUAGE, model_dir=str(root / "stanza"), processors=STANZA_PROCESSORS, verbose=False)
    data = {
        "schema_version": 1,
        "reranker": {"model": RERANKER_MODEL, "revision": RERANKER_REVISION},
        "stanza": {"version": stanza.__version__, "language": STANZA_LANGUAGE, "processors": STANZA_PROCESSORS},
        "files": _hashes(root, ["reranker", "stanza"]),
    }
    (root / "extras.json").write_text(json.dumps(data, sort_keys=True, indent=2) + "\n")
    extras_probe(root)
    return data


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("command", choices=["prepare", "verify", "prepare-extras", "verify-extras"])
    parser.add_argument("--output", type=Path, default=asset_root())
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--revision")
    parser.add_argument("--encoding", action="append")
    args = parser.parse_args()
    if args.command in {"prepare-extras", "verify-extras"}:
        data = prepare_extras(args.output) if args.command == "prepare-extras" else verify_extras(args.output)
        if args.command == "verify-extras":
            extras_probe(args.output)
        print(json.dumps({"reranker": data["reranker"], "stanza": data["stanza"], "files": len(data["files"])}))
        return
    if args.command == "prepare":
        data = prepare(
            args.output,
            args.model,
            args.revision,
            list(dict.fromkeys(args.encoding or ["cl100k_base"])),
        )
    else:
        data = verify(args.output)
        offline_probe(args.output)
    print(
        json.dumps(
            {
                "model": data["embedding"],
                "encodings": data["encodings"],
                "files": len(data["files"]),
            }
        )
    )


if __name__ == "__main__":
    main()