so#!/usr/bin/env python3
"""
Setup script to create a virtual environment and install dependencies for the Prerad project.
This script creates a venv at the project path and installs all necessary packages.
"""

import os
import sys
import subprocess
import venv
from pathlib import Path

def main():
    # Define paths
    project_path = Path("/run/media/dheena/Leave you files/prerad")
    venv_path = project_path / "venv"
    
    print(f"🚀 Setting up Prerad project environment")
    print(f"📁 Project path: {project_path}")
    print(f"📦 Virtual environment path: {venv_path}")
    print()
    
    # Step 1: Create virtual environment
    print("Step 1: Creating virtual environment...")
    try:
        venv.create(venv_path, with_pip=True)
        print(f"✅ Virtual environment created at {venv_path}")
    except Exception as e:
        print(f"❌ Error creating virtual environment: {e}")
        sys.exit(1)
    
    # Step 2: Determine pip executable path
    if sys.platform == "win32":
        pip_exe = venv_path / "Scripts" / "pip.exe"
        python_exe = venv_path / "Scripts" / "python.exe"
    else:
        pip_exe = venv_path / "bin" / "pip"
        python_exe = venv_path / "bin" / "python"
    
    print()
    print("Step 2: Upgrading pip...")
    try:
        subprocess.check_call([str(pip_exe), "install", "--upgrade", "pip"])
        print("✅ Pip upgraded successfully")
    except Exception as e:
        print(f"❌ Error upgrading pip: {e}")
        sys.exit(1)
    
    # Step 3: Install requirements for Streamlit (UI application)
    print()
    print("Step 3: Installing Streamlit and dependencies...")
    streamlit_requirements = [
        "streamlit>=1.28.0",
        "transformers>=4.36.0",
        "pillow>=10.0.0",
        "torch>=2.0.0",  # CPU-only PyTorch
    ]
    
    # Install torch from CPU index first
    print("  - Installing PyTorch (CPU version)...")
    try:
        subprocess.check_call([
            str(pip_exe), "install", 
            "--index-url", "https://download.pytorch.org/whl/cpu",
            "torch>=2.0.0"
        ])
        print("  ✅ PyTorch installed")
    except Exception as e:
        print(f"  ⚠️  Warning installing PyTorch: {e}")
    
    # Install other streamlit packages
    print("  - Installing other packages...")
    try:
        subprocess.check_call([str(pip_exe), "install"] + streamlit_requirements[:-1])
        print("✅ Streamlit and dependencies installed")
    except Exception as e:
        print(f"❌ Error installing Streamlit dependencies: {e}")
        sys.exit(1)
    
    # Step 4: Install ETL dependencies
    print()
    print("Step 4: Installing ETL dependencies...")
    etl_requirements = [
        "jsonlines>=4.0.0",
        "pandas>=2.0.0",
    ]
    try:
        subprocess.check_call([str(pip_exe), "install"] + etl_requirements)
        print("✅ ETL dependencies installed")
    except Exception as e:
        print(f"❌ Error installing ETL dependencies: {e}")
        sys.exit(1)
    
    # Step 5: Install Jupyter
    print()
    print("Step 5: Installing Jupyter...")
    try:
        subprocess.check_call([str(pip_exe), "install", "jupyter>=1.0.0"])
        print("✅ Jupyter installed")
    except Exception as e:
        print(f"❌ Error installing Jupyter: {e}")
        sys.exit(1)
    
    # Step 6: Create activation scripts info
    print()
    print("=" * 60)
    print("✅ Setup Complete!")
    print("=" * 60)
    print()
    print("📌 To activate the virtual environment, run:")
    print()
    if sys.platform == "win32":
        print(f"   {venv_path}\\Scripts\\activate")
    else:
        print(f"   source {venv_path}/bin/activate")
    print()
    print("🚀 To start the Streamlit application after activation:")
    print(f"   streamlit run {project_path}/containers/streamlit/app.py")
    print()
    print("📓 To start Jupyter Notebook after activation:")
    print(f"   jupyter notebook --notebook-dir={project_path}/volumes/notebooks")
    print()
    print("🔧 Or use Docker Compose for containerized services:")
    print(f"   cd {project_path}")
    print("   sudo docker compose up -d jupyter streamlit")
    print()

if __name__ == "__main__":
    main()