import os
import sys
import argparse
import subprocess
import venv

# Force line buffering
sys.stdout.reconfigure(line_buffering=True)

def run_command(command, description, required=True):
    print(f"[STATUS] {description}...")
    try:
        process = subprocess.Popen(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            shell=True,
            bufsize=1
        )

        for line in process.stdout:
            line_str = line.strip()
            if "Downloading" in line_str or "Installing" in line_str:
                print(f"[STATUS] {line_str}")
            elif line_str:
                print(f"[INFO] {line_str}")

        process.wait()
        if process.returncode != 0:
            print(f"[ERROR] {description} failed (exit code {process.returncode})", file=sys.stderr)
            if required:
                return False
        return True
    except Exception as e:
        print(f"[ERROR] Exception during {description}: {str(e)}", file=sys.stderr)
        return not required

def verify_install(python_path):
    print("[STATUS] Verifying PyTorch installation...")
    verify_cmd = (
        f'"{python_path}" -c "import torch; import transformers; import PIL; '
        f'print(f\\"torch={{torch.__version__}}\\")"'
    )
    return run_command(verify_cmd, "Verifying installed packages", required=True)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data-dir', help='Writable directory for .venv (used by packaged app builds)')
    args = parser.parse_args()

    backend_dir = os.path.dirname(os.path.abspath(__file__))
    base_dir = os.path.abspath(args.data_dir) if args.data_dir else os.path.dirname(backend_dir)
    venv_dir = os.path.join(base_dir, ".venv")
    requirements_path = os.path.join(backend_dir, "requirements.txt")

    if not os.path.exists(venv_dir):
        print("[STATUS] Creating Python virtual environment...")
        try:
            venv.create(venv_dir, with_pip=True)
            print("[STATUS] Virtual environment successfully created.")
        except Exception as e:
            print(f"[ERROR] Failed to create virtual environment: {str(e)}", file=sys.stderr)
            sys.exit(1)
    else:
        print("[STATUS] Virtual environment already exists.")

    if os.name == 'nt':
        python_path = os.path.join(venv_dir, "Scripts", "python.exe")
    else:
        python_path = os.path.join(venv_dir, "bin", "python")

    if not os.path.exists(python_path):
        print("[ERROR] Python executable not found inside virtual environment!", file=sys.stderr)
        sys.exit(1)

    run_command(
        f'"{python_path}" -m pip install --upgrade pip',
        "Upgrading pip to latest version",
        required=False
    )

    torch_index = "https://download.pytorch.org/whl/cpu"
    torch_cmd = (
        f'"{python_path}" -m pip install torch torchvision '
        f'--index-url {torch_index}'
    )
    if not run_command(torch_cmd, "Installing PyTorch AI Engine (CPU version)"):
        print("[ERROR] PyTorch installation failed. Ensure Python 3.10–3.12 is installed and you have internet access.", file=sys.stderr)
        sys.exit(1)

    other_deps_cmd = (
        f'"{python_path}" -m pip install '
        f'"transformers<5" timm einops kornia pillow accelerate'
    )
    if not run_command(other_deps_cmd, "Installing HuggingFace Transformers & PIL library"):
        print("[ERROR] Dependency installation failed!", file=sys.stderr)
        sys.exit(1)

    if not verify_install(python_path):
        print("[ERROR] Package verification failed after installation!", file=sys.stderr)
        sys.exit(1)

    print("[STATUS] Setup completed successfully!")
    print("[STATUS] DONE")

if __name__ == '__main__':
    main()