import json
import shutil
import tarfile
from pathlib import Path


BASE_DIR = Path(__file__).resolve().parent
SRC_TAR = BASE_DIR / "model" / "20260529-092118-humid-stool.tar.gz"
WORK_DIR = BASE_DIR / "evil_import_model_expanded"
OUT_TAR = BASE_DIR / "evil_import_model.tar.gz"
FALLBACK_OUT_TAR = BASE_DIR / "evil_import_model.new.tar.gz"


def main() -> None:
    if not SRC_TAR.is_file():
        raise SystemExit(f"missing source archive: {SRC_TAR}")

    if WORK_DIR.exists():
        shutil.rmtree(WORK_DIR)
    WORK_DIR.mkdir(parents=True)

    with tarfile.open(SRC_TAR, "r:gz") as tar:
        tar.extractall(WORK_DIR, filter="data")

    metadata_path = WORK_DIR / "metadata.json"
    if not metadata_path.is_file():
        raise SystemExit(f"missing metadata.json: {metadata_path}")

    metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
    metadata["predict_schema"]["nodes"]["pwn_import"] = {
        "needs": {},
        "uses": "pwnmod.Pwn",
        "constructor_name": "load",
        "fn": "provide",
        "config": {},
        "eager": True,
        "is_target": False,
        "is_input": False,
        "resource": None,
    }
    metadata_path.write_text(
        json.dumps(metadata, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )

    output_path = OUT_TAR
    if OUT_TAR.exists():
        try:
            OUT_TAR.unlink()
        except PermissionError:
            output_path = FALLBACK_OUT_TAR

    if output_path.exists():
        output_path.unlink()

    with tarfile.open(output_path, "w:gz") as tar:
        for item in WORK_DIR.iterdir():
            tar.add(item, arcname=item.name)

    print(f"source: {SRC_TAR}")
    print(f"output: {output_path}")


if __name__ == "__main__":
    main()