import os
import sys
import subprocess
import venv

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

def run_command(command, description):
    print(f"[STATUS] {description}...")
    try:
        process = subprocess.Popen(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            shell=True,
            bufsize=1
        )
        
        # Read output in real-time
        for line in process.stdout:
            line_str = line.strip()
            # If the output has pip download status, keep it readable but don't spam too much
            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] Command failed with code {process.returncode}", file=sys.stderr)
            return False
        return True
    except Exception as e:
        print(f"[ERROR] Exception occurred: {str(e)}", file=sys.stderr)
        return False

def main():
    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    venv_dir = os.path.join(base_dir, ".venv")
    
    # 1. Create venv if not exists
    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.")

    # 2. Path to pip inside the venv
    if os.name == 'nt':
        pip_path = os.path.join(venv_dir, "Scripts", "pip.exe")
        python_path = os.path.join(venv_dir, "Scripts", "python.exe")
    else:
        pip_path = os.path.join(venv_dir, "bin", "pip")
        python_path = os.path.join(venv_dir, "bin", "python")

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

    # 3. Upgrade pip
    run_command(f'"{pip_path}" install --upgrade pip', "Upgrading pip to latest version")

    # 4. Install dependencies
    requirements_path = os.path.join(base_dir, "backend", "requirements.txt")
    
    # Use extra index URL for CPU-only torch to download faster (150MB instead of 2.5GB+ GPU version)
    # This is perfect for single images on desktop!
    pip_install_cmd = f'"{pip_path}" install -r "{requirements_path}" --extra-index-url https://download.pytorch.org/whl/cpu'
    
    success = run_command(pip_install_cmd, "Installing PyTorch, Transformers, and PIL dependencies")
    if not success:
        print("[ERROR] Dependency installation failed!", file=sys.stderr)
        sys.exit(1)

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

if __name__ == '__main__':
    main()