# This script sets up the entire local project environment on Windows.
# It creates folders, creates a virtual environment, upgrades pip tools,
# and installs all dependencies from requirements.txt.

$ErrorActionPreference = "Stop"

function Invoke-StepCommand {
    param(
        [string]$Description,
        [scriptblock]$CommandBlock
    )
    Write-Host $Description -ForegroundColor Cyan
    & $CommandBlock
    if ($LASTEXITCODE -ne 0) {
        throw "Command failed during: $Description"
    }
}

try {
    # This finds the project root based on this script location.
    $scriptDir = Split-Path -Parent $MyInvocation.MyCommand.Path
    $projectRoot = Split-Path -Parent $scriptDir
    Set-Location $projectRoot

    Write-Host "Step 1/6: Creating project folder structure..." -ForegroundColor Cyan

    # These are all folders needed for Components 1-10.
    $folders = @(
        "data",
        "data/raw",
        "data/interim",
        "data/processed",
        "data/external",
        "src",
        "src/tokenizer",
        "src/dataset_pipeline",
        "src/model_architecture",
        "src/training_pipeline",
        "src/evaluation_system",
        "src/inference_engine",
        "src/chat_interface",
        "src/finetuning_system",
        "src/export_optimization",
        "configs",
        "scripts",
        "tests",
        "checkpoints",
        "models",
        "models/base",
        "models/lora",
        "models/quantized",
        "artifacts",
        "logs"
    )

    foreach ($folder in $folders) {
        if (-not (Test-Path $folder)) {
            New-Item -ItemType Directory -Path $folder | Out-Null
        }
    }

    Write-Host "Step 2/6: Creating Python virtual environment..." -ForegroundColor Cyan

    # This checks Python version before environment creation.
    # We require Python 3.10 or 3.11 for best Windows compatibility with ML packages.
    $pyVersionRaw = python -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')"
    if ($LASTEXITCODE -ne 0) {
        throw "Python is not available in PATH. Install Python 3.11 and try again."
    }
    if (($pyVersionRaw -ne "3.10") -and ($pyVersionRaw -ne "3.11")) {
        throw "Detected Python $pyVersionRaw. Please install Python 3.11 (recommended) or 3.10, then rerun this script."
    }

    # This creates .venv only if it does not already exist.
    if (-not (Test-Path ".venv")) {
        python -m venv .venv
    } else {
        Write-Host "Virtual environment already exists. Reusing .venv." -ForegroundColor Yellow
    }

    # This points to the venv Python executable on Windows.
    $venvPython = Join-Path $projectRoot ".venv\Scripts\python.exe"
    if (-not (Test-Path $venvPython)) {
        throw "Could not find .venv Python at $venvPython. Please check Python installation."
    }

    Invoke-StepCommand "Step 3/6: Upgrading pip, setuptools, and wheel..." {
        & $venvPython -m pip install --upgrade pip setuptools wheel
    }

    Invoke-StepCommand "Step 4/6: Installing CUDA-enabled PyTorch (cu121)..." {
        & $venvPython -m pip install --index-url https://download.pytorch.org/whl/cu121 torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1
    }

    Invoke-StepCommand "Step 5/6: Installing project dependencies from requirements.txt..." {
        & $venvPython -m pip install -r requirements.txt
    }

    Invoke-StepCommand "Step 6/6: Validating pip environment health..." {
        & $venvPython -m pip check
    }

    Write-Host "Setup complete." -ForegroundColor Green
    Write-Host "Next command: .\.venv\Scripts\Activate.ps1" -ForegroundColor Green
    Write-Host "Then run: python .\scripts\verify_component1_setup.py" -ForegroundColor Green
}
catch {
    # This prints a clear plain-English error if anything breaks.
    Write-Host ""
    Write-Host "Setup failed." -ForegroundColor Red
    Write-Host "What went wrong: $($_.Exception.Message)" -ForegroundColor Red
    Write-Host "Fix suggestion: Check Python is installed and available in PATH, then run this script again." -ForegroundColor Yellow
    exit 1
}