#!/usr/bin/env python3

import os
import re
import shutil
import sys
import argparse


def verify_venv(venv):
    """Verify that the virtual environment is valid"""
    if not os.path.isdir(venv):
        print("The virtual environment {} does not exist".format(venv))
        return False
    if not os.path.isdir(os.path.join(venv, "bin")):
        print("The virtual environment {} does not contain a bin directory".format(venv))
        return False
    if not os.path.isfile(os.path.join(venv, "bin", "activate")):
        print("The virtual environment {} does not contain a bin/activate file".format(venv))
        return False
    return True

def remove_pycache_dirs(root: str) -> None:
    # Remove existing __pycache__ directories
    for dirpath, dirs, files in os.walk(root):
        for d in dirs:
            if d == "__pycache__":
                shutil.rmtree(os.path.join(dirpath, d))

def venv_move(venv: str, remove_pycache: bool=True, yes: bool=False) -> None:
    if not verify_venv(venv):
        raise ValueError("Invalid virtual environment")

    venv = os.path.abspath(venv)
    venv_name = os.path.basename(venv)

    if remove_pycache:
        remove_pycache_dirs(venv)

    # Find and replace old path with new path
    activate_path = os.path.join(venv, "bin", "activate")
    with open(activate_path, "r") as f:
        activate_script = f.read()

    old_path = re.search('VIRTUAL_ENV="(.*?)"', activate_script).group(1)
    new_path = os.path.join(os.getcwd(), venv_name)

    if old_path == new_path:
        print("venv paths are already set correctly to {}".format(new_path))
        return

    matching_files = []

    for root, dirs, files in os.walk(venv): for file in files:
        ext = os.path.splitext(file)[1]
        if ext not in [".py", ".sh"]:
            continue
        file_path = os.path.join(root, file)
        with open(file_path, "r") as f:
            file_text = f.read()
        if old_path in file_text:
            print(file_path)
            matching_files.append(file_path)

    if not yes:
        response = input("Replace {} with {} in the above files? [y/n]".format(old_path, new_path))
    if not(yes or response == "y"):
        return

    for file_path in matching_files:
        with open(file_path, "r") as f:
            file_text = f.read()
        file_text = file_text.replace(old_path, new_path)
        with open(file_path, "w") as f:
            f.write(file_text)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("venv", help="Path to the virtual environment")
    parser.add_argument("--keep-pycache", action="store_true", help="Do not remove __pycache__ directories")
    parser.add_argument("--yes", action="store_true", help="Do not prompt for confirmation")
    args = parser.parse_args()
    venv_move(args.venv, remove_pycache=not args.keep_pycache, yes=args.yes)