#!/usr/bin/env python3
"""
Verification script to check if the project is ready for Zero GPU deployment
"""

import os
import sys
import re

def check_file_exists(filepath, required=True):
    """Check if a file exists"""
    exists = os.path.exists(filepath)
    status = "✅" if exists else ("❌" if required else "⚠️")
    print(f"{status} {filepath}: {'Found' if exists else 'Missing'}")
    return exists

def check_pytorch_version():
    """Check PyTorch version in requirements.txt"""
    print("\n📦 Checking PyTorch version...")

    with open("requirements.txt", "r") as f:
        content = f.read()

    # Look for torch version
    torch_match = re.search(r"torch==(\d+\.\d+\.\d+)", content)
    if torch_match:
        version = torch_match.group(1)
        if version == "2.2.0":
            print(f"✅ PyTorch version {version} is compatible with Zero GPU")
            return True
        else:
            print(f"❌ PyTorch version {version} may not be compatible with Zero GPU")
            print("   Recommended: torch==2.2.0")
            return False
    else:
        print("❌ PyTorch version not found in requirements.txt")
        return False

def check_spaces_not_in_requirements():
    """Check that spaces is NOT in requirements.txt"""
    print("\n📦 Checking spaces library...")

    with open("requirements.txt", "r") as f:
        content = f.read()

    if "spaces" in content:
        print("❌ 'spaces' found in requirements.txt - should be removed!")
        print("   HF Spaces provides this automatically")
        return False
    else:
        print("✅ 'spaces' correctly not in requirements.txt")
        return True

def check_gpu_decorators():
    """Check for @spaces.GPU decorators"""
    print("\n🎯 Checking GPU decorators...")

    files_to_check = ["mvp.py"]
    decorator_pattern = r"@spaces\.GPU"

    for filepath in files_to_check:
        if not os.path.exists(filepath):
            print(f"⚠️ {filepath} not found")
            continue

        with open(filepath, "r") as f:
            content = f.read()

        decorators = re.findall(decorator_pattern, content)
        if decorators:
            print(f"✅ {filepath}: Found {len(decorators)} @spaces.GPU decorators")

            # Check specific functions
            functions = ["run_model", "reconstruct", "detect_objects"]
            for func in functions:
                # Check if function has decorator
                pattern = rf"@spaces\.GPU.*?\ndef {func}\("
                if re.search(pattern, content, re.DOTALL):
                    print(f"   ✅ {func}() has GPU decorator")
                else:
                    print(f"   ⚠️ {func}() might be missing GPU decorator")
        else:
            print(f"❌ {filepath}: No GPU decorators found")
            return False

    return True

def check_model_manager():
    """Check for ModelManager implementation"""
    print("\n🔧 Checking ModelManager...")

    with open("mvp.py", "r") as f:
        content = f.read()

    if "class ModelManager" in content:
        print("✅ ModelManager class found")

        # Check for key methods
        methods = ["get_vggt_model", "get_metric3d_model", "get_clip_model", "clear_cache"]
        for method in methods:
            if f"def {method}" in content:
                print(f"   ✅ {method}() method found")
            else:
                print(f"   ❌ {method}() method missing")

        # Check for global instance
        if "model_manager = ModelManager()" in content:
            print("   ✅ Global model_manager instance created")
        else:
            print("   ⚠️ Global model_manager instance not found")

        return True
    else:
        print("❌ ModelManager class not found")
        return False

def check_environment_variables():
    """Check for Zero GPU environment variable handling"""
    print("\n🌍 Checking environment variable handling...")

    files = ["app.py", "mvp.py"]

    for filepath in files:
        with open(filepath, "r") as f:
            content = f.read()

        checks = {
            "IS_HF_SPACE": "HF Space detection",
            "IS_ZERO_GPU": "Zero GPU mode detection",
            "MAX_IMAGES": "Image limit configuration",
            "BATCH_SIZE": "Batch size configuration",
        }

        print(f"\n  {filepath}:")
        for var, description in checks.items():
            if var in content:
                print(f"   ✅ {var}: {description}")
            else:
                print(f"   ⚠️ {var}: {description} not found")

def check_readme_header():
    """Check README header for HF Spaces"""
    print("\n📄 Checking README configuration...")

    if os.path.exists("README_HF.md"):
        with open("README_HF.md", "r") as f:
            content = f.read()

        if content.startswith("---"):
            # Extract YAML header
            header_end = content.find("---", 3)
            if header_end > 0:
                header = content[3:header_end]

                required_fields = {
                    "title:": "Title field",
                    "emoji:": "Emoji field",
                    "sdk: gradio": "Gradio SDK",
                    "app_file: app.py": "App file specification",
                }

                print("✅ README_HF.md has YAML header")
                for field, description in required_fields.items():
                    if field in header:
                        print(f"   ✅ {description}")
                    else:
                        print(f"   ❌ {description} missing")
            else:
                print("❌ README_HF.md YAML header not properly closed")
        else:
            print("❌ README_HF.md missing YAML header")
    else:
        print("❌ README_HF.md not found")

def main():
    """Run all verification checks"""
    print("=" * 60)
    print("🔍 Zoo3D Zero GPU Deployment Verification")
    print("=" * 60)

    # Check critical files
    print("\n📁 Checking required files...")
    check_file_exists("app.py")
    check_file_exists("mvp.py")
    check_file_exists("requirements.txt")
    check_file_exists("packages.txt")
    check_file_exists("README_HF.md")
    check_file_exists("ZERO_GPU_README.md", required=False)
    check_file_exists("test_zero_gpu.py", required=False)

    # Check configurations
    pytorch_ok = check_pytorch_version()
    spaces_ok = check_spaces_not_in_requirements()
    decorators_ok = check_gpu_decorators()
    model_manager_ok = check_model_manager()
    check_environment_variables()
    check_readme_header()

    # Summary
    print("\n" + "=" * 60)
    print("📊 Summary")
    print("=" * 60)

    critical_checks = [pytorch_ok, spaces_ok, decorators_ok, model_manager_ok]

    if all(critical_checks):
        print("✅ Project is ready for Zero GPU deployment!")
        print("\nNext steps:")
        print("1. Copy README_HF.md to README.md when deploying")
        print("2. Push to Hugging Face Spaces with Zero GPU hardware")
        print("3. Monitor logs for any runtime issues")
        return 0
    else:
        print("❌ Some issues need to be fixed before deployment")
        print("\nCritical issues to resolve:")
        if not pytorch_ok:
            print("- Update PyTorch to version 2.2.0")
        if not spaces_ok:
            print("- Remove 'spaces' from requirements.txt")
        if not decorators_ok:
            print("- Add @spaces.GPU decorators to GPU functions")
        if not model_manager_ok:
            print("- Implement ModelManager for lazy loading")
        return 1

if __name__ == "__main__":
    sys.exit(main())