#!/usr/bin/env python3
"""
Startup Fix Script for Dressify
Handles dataset preparation issues and ensures system startup
"""

import os
import sys
import subprocess
import time
from pathlib import Path


def check_dataset_status():
    """Check the current dataset status."""
    print("🔍 Checking dataset status...")
    
    root = os.path.abspath(os.path.join(os.getcwd(), "data", "Polyvore"))
    
    if not os.path.exists(root):
        print(f"❌ Dataset directory not found: {root}")
        return False
    
    # Check key components
    images_dir = os.path.join(root, "images")
    splits_dir = os.path.join(root, "splits")
    
    has_images = os.path.isdir(images_dir) and any(Path(images_dir).glob("*"))
    has_splits = os.path.isdir(splits_dir) and any(Path(splits_dir).glob("*.json"))
    
    print(f"📁 Dataset root: {root}")
    print(f"🖼️ Images: {'✅' if has_images else '❌'} ({images_dir})")
    print(f"📊 Splits: {'✅' if has_splits else '❌'} ({splits_dir})")
    
    # Check for official splits
    official_splits = []
    for location in ["nondisjoint", "disjoint"]:
        location_path = os.path.join(root, location)
        if os.path.exists(location_path):
            for split in ["train", "valid", "test"]:
                split_file = os.path.join(location_path, f"{split}.json")
                if os.path.exists(split_file):
                    size_mb = os.path.getsize(split_file) / (1024 * 1024)
                    official_splits.append(f"{location}/{split}.json ({size_mb:.1f} MB)")
    
    if official_splits:
        print(f"🎯 Official splits found:")
        for split in official_splits:
            print(f"   ✅ {split}")
    
    if has_images and has_splits:
        print("✅ Dataset is ready!")
        return True
    elif has_images:
        print("⚠️ Images present but splits missing - will create splits from official data")
        return "needs_splits"
    else:
        print("❌ Dataset incomplete - needs full preparation")
        return False


def prepare_dataset():
    """Prepare the dataset using the improved scripts."""
    print("\n🚀 Preparing dataset...")
    
    root = os.path.abspath(os.path.join(os.getcwd(), "data", "Polyvore"))
    
    # First, ensure the data fetcher runs
    try:
        print("📥 Running data fetcher...")
        from utils.data_fetch import ensure_dataset_ready
        dataset_root = ensure_dataset_ready()
        
        if not dataset_root:
            print("❌ Data fetcher failed")
            return False
            
        print(f"✅ Data fetcher completed: {dataset_root}")
        
    except Exception as e:
        print(f"❌ Data fetcher error: {e}")
        return False
    
    # Now run the dataset preparation script (without random splits)
    try:
        print("🔧 Running dataset preparation...")
        
        # Check if prepare_polyvore.py exists
        prep_script = "scripts/prepare_polyvore.py"
        if not os.path.exists(prep_script):
            prep_script = "prepare_polyvore.py"
        
        if not os.path.exists(prep_script):
            print(f"❌ Prepare script not found: {prep_script}")
            return False
        
        # Run the preparation script WITHOUT random splits
        cmd = [
            sys.executable, prep_script,
            "--root", root
            # Note: NOT using --force_random_split
        ]
        
        print(f"🔧 Running: {' '.join(cmd)}")
        print("🎯 This will use official splits from nondisjoint/ and disjoint/ folders")
        
        result = subprocess.run(cmd, capture_output=True, text=True, check=False)
        
        if result.returncode == 0:
            print("✅ Dataset preparation completed successfully!")
            print("📝 Output:")
            print(result.stdout)
            return True
        else:
            print("❌ Dataset preparation failed!")
            print("📝 Error output:")
            print(result.stderr)
            print("📝 Standard output:")
            print(result.stdout)
            
            # Check if it's because official splits are missing
            if "No official splits found" in result.stderr or "No official splits found" in result.stdout:
                print("\n🔧 Issue: Official splits not found in nondisjoint/ or disjoint/ folders")
                print("📁 Expected structure:")
                print("   data/Polyvore/")
                print("   ├── nondisjoint/")
                print("   │   ├── train.json")
                print("   │   ├── valid.json")
                print("   │   └── test.json")
                print("   ├── disjoint/")
                print("   │   ├── train.json")
                print("   │   ├── valid.json")
                print("   │   └── test.json")
                print("   └── images/")
                
                print("\n💡 Solution: The dataset should have been downloaded with official splits.")
                print("   Check if the Hugging Face download completed successfully.")
                
            return False
            
    except Exception as e:
        print(f"❌ Dataset preparation error: {e}")
        return False


def verify_splits():
    """Verify that splits were created successfully."""
    print("\n🔍 Verifying splits...")
    
    root = os.path.abspath(os.path.join(os.getcwd(), "data", "Polyvore"))
    splits_dir = os.path.join(root, "splits")
    
    if not os.path.exists(splits_dir):
        print("❌ Splits directory not found")
        return False
    
    required_files = [
        "train.json",
        "outfits_train.json", 
        "outfit_triplets_train.json"
    ]
    
    missing_files = []
    for file_name in required_files:
        file_path = os.path.join(splits_dir, file_name)
        if os.path.exists(file_path):
            size_mb = os.path.getsize(file_path) / (1024 * 1024)
            print(f"✅ {file_name}: {size_mb:.1f} MB")
        else:
            print(f"❌ {file_name}: Missing")
            missing_files.append(file_name)
    
    if missing_files:
        print(f"❌ Missing required files: {missing_files}")
        return False
    
    print("✅ All required splits verified!")
    return True


def test_training_scripts():
    """Test that training scripts can run without errors."""
    print("\n🧪 Testing training scripts...")
    
    # Test ResNet training script
    try:
        print("🔧 Testing ResNet training script...")
        from models.resnet_embedder import ResNetItemEmbedder
        print("✅ ResNet model imports successfully")
    except Exception as e:
        print(f"❌ ResNet model import failed: {e}")
        return False
    
    # Test ViT training script
    try:
        print("🔧 Testing ViT training script...")
        from models.vit_outfit import OutfitCompatibilityModel
        print("✅ ViT model imports successfully")
    except Exception as e:
        print(f"❌ ViT model import failed: {e}")
        return False
    
    print("✅ All training scripts tested successfully!")
    return True


def create_quick_start_script():
    """Create a quick start script for easy testing."""
    script_content = """#!/bin/bash
# Quick Start Script for Dressify
# This script will prepare the dataset and start training

echo "🚀 Dressify Quick Start"
echo "========================"

# Check if dataset is ready
if [ -d "data/Polyvore/splits" ] && [ -f "data/Polyvore/splits/train.json" ]; then
    echo "✅ Dataset is ready!"
else
    echo "🔧 Preparing dataset..."
    python startup_fix.py
fi

# Start quick training
echo "🎯 Starting quick training..."
python train_resnet.py --data_root data/Polyvore --epochs 3 --out models/exports/resnet_quick.pth

echo "🎉 Quick start completed!"
echo "📁 Check models/exports/ for trained models"
"""
    
    script_path = "quick_start.sh"
    with open(script_path, "w") as f:
        f.write(script_content)
    
    # Make executable
    os.chmod(script_path, 0o755)
    print(f"📝 Created quick start script: {script_path}")


def main():
    """Main startup fix routine."""
    print("🚀 Dressify Startup Fix")
    print("=" * 50)
    
    # Check current status
    status = check_dataset_status()
    
    if status is True:
        print("✅ System is ready to go!")
        return True
    
    elif status == "needs_splits":
        print("🔧 Dataset needs splits created from official data...")
        if prepare_dataset():
            if verify_splits():
                print("✅ Dataset preparation completed successfully!")
                return True
            else:
                print("❌ Split verification failed")
                return False
        else:
            print("❌ Dataset preparation failed")
            return False
    
    else:
        print("🔧 Dataset needs full preparation...")
        if prepare_dataset():
            if verify_splits():
                print("✅ Dataset preparation completed successfully!")
                return True
            else:
                print("❌ Split verification failed")
                return False
        else:
            print("❌ Dataset preparation failed")
            return False


if __name__ == "__main__":
    try:
        success = main()
        
        if success:
            print("\n🎉 Startup fix completed successfully!")
            print("🚀 Your Dressify system is ready to use!")
            
            # Create quick start script
            create_quick_start_script()
            
            print("\n📋 Next steps:")
            print("1. Run: python app.py")
            print("2. Or use: ./quick_start.sh")
            print("3. Check the Advanced Training tab for parameter controls")
            
        else:
            print("\n❌ Startup fix failed!")
            print("🔧 Please check the error messages above")
            print("📞 Contact support if issues persist")
            
    except KeyboardInterrupt:
        print("\n⏹️ Startup fix interrupted by user")
    except Exception as e:
        print(f"\n💥 Unexpected error: {e}")
        import traceback
        traceback.print_exc()