#!/usr/bin/env python3
"""Patch pinned TransformerLab source to tolerate symlinked home directories."""

from __future__ import annotations

import argparse
import re
import sys
from pathlib import Path


PATCH_MARKER = "with symlinked TransformerLab home directories."
TARGET_BLOCK = re.compile(
    r"(?P<indent>[ \t]+)# The following prevents path traversal attacks:.*?"
    r"(?P=indent)# now get the file contents",
    re.DOTALL,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--transformerlab-dir", required=True)
    return parser.parse_args()


def repair_plugins_router(path: Path) -> bool:
    source = path.read_text(encoding="utf-8")
    if PATCH_MARKER in source:
        return False

    replacement = (
        "    # The following prevents path traversal attacks while remaining compatible\n"
        "    # with symlinked TransformerLab home directories.\n"
        "    plugin_dir = Path(await lab_dirs.plugin_dir_by_name((pluginId)))\n"
        "    resolved_plugin_dir = plugin_dir.resolve()\n"
        '    final_path = (plugin_dir / f"{filename}{file_ext}").resolve()\n'
        "\n"
        "    try:\n"
        "        final_path.relative_to(resolved_plugin_dir)\n"
        "    except ValueError:\n"
        '        return {"message": f"File {filename}{file_ext} is outside plugin directory"}\n'
        "\n"
        "    # now get the file contents"
    )
    updated, count = TARGET_BLOCK.subn(replacement, source, count=1)
    if count != 1:
        raise RuntimeError(f"Could not find path traversal block in {path}")

    path.write_text(updated, encoding="utf-8")
    return True


def main() -> int:
    args = parse_args()
    root = Path(args.transformerlab_dir).expanduser().resolve()
    plugins_router = root / "src" / "transformerlab" / "routers" / "experiment" / "plugins.py"
    if not plugins_router.exists():
        print(f"missing TransformerLab plugins router: {plugins_router}", file=sys.stderr)
        return 1

    changed = repair_plugins_router(plugins_router)
    if changed:
        print(f"patched {plugins_router}")
    else:
        print(f"already patched {plugins_router}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())