#!/usr/bin/env python3
"""
run_script_venv.py
==================
A utility that
1. Accepts a path to a Python script.
2. Creates an isolated working directory.
3. Creates a virtual environment inside that directory using **uv**.
4. Statically analyses the target script to discover third‑party dependencies.
5. Installs the detected dependencies into the environment.
6. Executes the target script inside the environment.

Usage
-----
$ python run_script_venv.py path/to/script.py [--workdir ./my_dir] [--keep-workdir]

Requirements
------------
* Python ≥ 3.10 (for ``sys.stdlib_module_names``)
* ``uv`` must be available on the host Python used to run this helper.

Notes
-----
* The dependency detection is based on import statements only. Runtime/optional
  dependencies that are conditionally imported or loaded by plugins are not
  detected; you can pass extra packages manually via ``--extra``.
* A small mapping translates common import‑name↔package‑name mismatches
  (e.g. ``cv2`` → ``opencv‑python``).
"""

from __future__ import annotations

import argparse
import ast
import os
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Iterable, Set

# --------------------------- Helpers -----------------------------------------

# Map import names to PyPI package names where they differ.
NAME_MAP: dict[str, str] = {
    "cv2": "opencv-python",
    "sklearn": "scikit-learn",
    "PIL": "pillow",
    "yaml": "pyyaml",
    "Crypto": "pycryptodome",
}


def find_imports(script_path: Path) -> Set[str]:
    """Return a set of *import names* found in *script_path*."""
    root = ast.parse(script_path.read_text())
    imports: set[str] = set()
    for node in ast.walk(root):
        if isinstance(node, ast.Import):
            for alias in node.names:
                imports.add(alias.name.split(".")[0])
        elif isinstance(node, ast.ImportFrom):
            if node.level == 0 and node.module:  # skip relative imports
                imports.add(node.module.split(".")[0])
    return imports


def third_party_modules(modules: Iterable[str]) -> Set[str]:
    """Filter *modules* to those not in the Python stdlib."""
    stdlib = set(sys.stdlib_module_names)  # Python ≥ 3.10
    return {m for m in modules if m not in stdlib}


def translate_names(modules: Iterable[str]) -> Set[str]:
    """Translate *modules* to their PyPI names using NAME_MAP."""
    pkgs: set[str] = set()
    for m in modules:
        pkgs.add(NAME_MAP.get(m, m))
    return pkgs


def create_venv(venv_dir: Path) -> Path:
    """Create a venv at *venv_dir* using ``uv venv`` and return the Python exe."""
    subprocess.check_call(["uv", "venv", str(venv_dir)])
    python_exe = venv_dir / ("Scripts" if os.name == "nt" else "bin") / "python"
    return python_exe


def install_packages(python_exe: Path, packages: Iterable[str]) -> None:
    if not packages:
        print("[+] No third‑party packages detected; skipping installation.")
        return
    print(f"[+] Installing: {' '.join(packages)}")
    # subprocess.check_call([str(python_exe), "-m", "pip", "install", *packages])
    subprocess.check_call(["uv", "pip", "install", "-p", str(python_exe), *packages])


# --------------------------- Main routine ------------------------------------

def main() -> None:
    p = argparse.ArgumentParser(prog="run_script_venv")
    p.add_argument("script", type=Path, help="Python script to execute")
    p.add_argument(
        "--workdir",
        type=Path,
        help="Directory to create the environment in (defaults to a temp dir)",
    )
    p.add_argument(
        "--extra",
        nargs="*",
        default=[],
        metavar="PKG",
        help="Extra PyPI packages to install in addition to auto‑detected ones",
    )
    p.add_argument(
        "--keep-workdir",
        action="store_true",
        help="Do not delete the temporary working directory on success",
    )

    args = p.parse_args()
    script_path: Path = args.script.resolve()
    if not script_path.is_file():
        p.error(f"Script '{script_path}' not found.")

    # ---------------------------------------------------------------------
    # Prepare working directory and venv
    # ---------------------------------------------------------------------
    workdir: Path
    temp_created = False
    if args.workdir:
        workdir = args.workdir.resolve()
        workdir.mkdir(parents=True, exist_ok=True)
    else:
        workdir = Path(tempfile.mkdtemp(prefix=f"{script_path.stem}_work_"))
        temp_created = True
    venv_dir = workdir / ".venv"
    print(f"[+] Working directory: {workdir}")

    # ---------------------------------------------------------------------
    # Detect and install dependencies
    # ---------------------------------------------------------------------
    mods = find_imports(script_path)
    third_party = translate_names(third_party_modules(mods)) | set(args.extra)

    print(f"[+] Third‑party imports: {sorted(third_party) if third_party else 'None'}")

    py = create_venv(venv_dir)
    install_packages(py, third_party)

    # ---------------------------------------------------------------------
    # Run the target script inside the environment
    # ---------------------------------------------------------------------
    env = os.environ.copy()
    env["VIRTUAL_ENV"] = str(venv_dir)
    env["PATH"] = str(py.parent) + os.pathsep + env["PATH"]

    print("[+] Running script …\n----------------------------------------")

    proc = subprocess.Popen(
        [str(py), str(script_path)],
        env=env,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
    )

    # Stream the script output line by line
    assert proc.stdout is not None
    for line in proc.stdout:
        print(line, end="")

    retcode = proc.wait()

    print("----------------------------------------")
    if retcode == 0:
        print("[+] Script finished successfully.")
    else:
        print(f"[!] Script exited with return code {retcode}.")
        sys.exit(retcode)


    # ---------------------------------------------------------------------
    # Cleanup
    # ---------------------------------------------------------------------
    if temp_created and not args.keep_workdir:
        shutil.rmtree(workdir)
        print(f"[+] Removed temporary directory {workdir}")


if __name__ == "__main__":
    try:
        main()
    except subprocess.CalledProcessError as e:
        print("[!] A subprocess exited with a non‑zero status:", e, file=sys.stderr)
        sys.exit(e.returncode)