"""Assemble the Hugging Face release folder: runtime code, documentation and the three adapters.

    python scripts/stage_release.py --out release \
        --operator DIR --guardian DIR --verifier DIR [--verifier-eval eval_summary.json]

Only the files listed here are copied. Training checkpoints, optimizer state, the generated training data and
PEFT's template README are left out. The script fails on anything that should not be published: a missing
adapter file, adapters trained on different base models, a local path in an adapter config, or an access token,
e-mail address or home-directory path in any text file.
"""
from __future__ import annotations

import argparse
import json
import re
import shutil
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
INCLUDE = ["README.md", "DEPLOY.md", "LICENSE", "requirements.txt", "requirements-gpu.txt", "sanad_text.py",
           "sanad_pipeline.py", "sanad_cases.py", "sanad_sim", "sanad_live", "training", "scripts", "tests",
           "deploy", "docs"]
SKIP = {"__pycache__", ".ipynb_checkpoints", ".DS_Store", ".pytest_cache"}
ADAPTER_FILES = ["adapter_config.json", "adapter_model.safetensors"]
TOKENIZER_FILES = ["tokenizer.json", "tokenizer_config.json", "chat_template.jinja", "special_tokens_map.json",
                   "added_tokens.json", "vocab.json", "merges.txt", "preprocessor_config.json",
                   "processor_config.json", "video_preprocessor_config.json"]
VERIFIER_PENDING = "| Verifier | Evaluation of this release is pending. | |"
FORBIDDEN = [(re.compile(r"\bhf_[A-Za-z0-9]{30,}"), "Hugging Face token"),
             (re.compile(r"[A-Za-z0-9._%+-]+@(?:gmail|outlook|hotmail|yahoo)\.com"), "e-mail address"),
             (re.compile(r"/home/[a-z]|/Users/[A-Za-z]|C:\\\\Users"), "home-directory path")]
TEXT = {".py", ".md", ".txt", ".json", ".ipynb", ".js", ".css", ".service", ".jinja"}


class ReleaseError(Exception):
    pass


def copy_code(out: Path) -> None:
    for name in INCLUDE:
        src = ROOT / name
        if not src.exists():
            raise ReleaseError(f"missing {name}")
        if src.is_dir():
            shutil.copytree(src, out / name, ignore=lambda d, names: [n for n in names if n in SKIP or n.endswith(".pyc")])
        else:
            shutil.copy2(src, out / name)


def copy_adapter(role: str, src: Path, out: Path) -> dict:
    dest = out / "adapters" / role
    dest.mkdir(parents=True)
    for f in ADAPTER_FILES:
        if not (src / f).is_file():
            raise ReleaseError(f"{role}: {f} not found in {src}")
        shutil.copy2(src / f, dest / f)
    if role == "operator":                         # one tokenizer and chat template, shared by all three adapters
        found = [f for f in TOKENIZER_FILES if (src / f).is_file()]
        if "tokenizer_config.json" not in found:
            raise ReleaseError(f"operator: tokenizer files not found in {src}")
        for f in found:
            shutil.copy2(src / f, dest / f)
    cfg = json.loads((dest / "adapter_config.json").read_text(encoding="utf-8"))
    base = cfg.get("base_model_name_or_path") or ""
    if not base or base.startswith(("/", ".")) or ":\\" in base:
        raise ReleaseError(f"{role}: base_model_name_or_path is a local path ({base!r})")
    return cfg


def fill_verifier_eval(out: Path, summary_file: Path | None) -> str:
    readme = out / "README.md"
    text = readme.read_text(encoding="utf-8")
    if VERIFIER_PENDING not in text:
        raise ReleaseError("README.md has no Verifier evaluation line")
    if summary_file is None:
        return "pending"
    s = json.loads(summary_file.read_text(encoding="utf-8"))
    line = (f"| Verifier | Status correct in {s['status_ok']}% of {s['scenarios']} held-out cases | "
            f"Missed failures: {s['missed_failure']}% of cases that were not resolved |")
    readme.write_text(text.replace(VERIFIER_PENDING, line), encoding="utf-8")
    return line


def scan(out: Path) -> list[str]:
    problems = []
    for p in sorted(out.rglob("*")):
        if p.is_file() and p.suffix in TEXT:
            text = p.read_text(encoding="utf-8", errors="replace")
            for rx, what in FORBIDDEN:
                for m in rx.finditer(text):
                    problems.append(f"{p.relative_to(out)}: {what}: {m.group(0)[:40]}")
    return problems


def stage(out: Path, adapters: dict[str, Path], verifier_eval: Path | None = None) -> dict:
    if out.exists():
        shutil.rmtree(out)
    out.mkdir(parents=True)
    copy_code(out)
    cfgs = {role: copy_adapter(role, Path(path), out) for role, path in adapters.items()}
    bases = {c["base_model_name_or_path"] for c in cfgs.values()}
    if len(bases) != 1:
        raise ReleaseError(f"the adapters were trained on different base models: {sorted(bases)}")
    verifier = fill_verifier_eval(out, verifier_eval)
    problems = scan(out)
    if problems:
        raise ReleaseError("not publishable:\n  " + "\n  ".join(problems))
    files = sorted(p for p in out.rglob("*") if p.is_file())
    return {"files": len(files), "bytes": sum(p.stat().st_size for p in files), "base_model": bases.pop(),
            "lora": {r: {"r": c.get("r"), "alpha": c.get("lora_alpha")} for r, c in cfgs.items()},
            "verifier_eval": verifier,
            "adapters": {str(p.relative_to(out)): p.stat().st_size for p in files
                         if p.relative_to(out).parts[0] == "adapters"}}


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--out", required=True, type=Path)
    for role in ("operator", "guardian", "verifier"):
        ap.add_argument(f"--{role}", required=True, type=Path, help=f"{role} adapter folder")
    ap.add_argument("--verifier-eval", type=Path, help="eval_summary.json of the Verifier run being published")
    a = ap.parse_args()
    try:
        info = stage(a.out, {"operator": a.operator, "guardian": a.guardian, "verifier": a.verifier}, a.verifier_eval)
    except ReleaseError as e:
        sys.exit(f"error: {e}")
    print(json.dumps(info, indent=1))


if __name__ == "__main__":
    main()