import os
import sys
import importlib.util
import subprocess

# Define required directories
required_dirs = [
    "omni_speech",
    "omni_speech/serve",
    "omni_speech/infer",
    "vocoder"
]

# Define required files
required_files = [
    "app.py",
    "omni_speech/__init__.py",
    "omni_speech/serve/__init__.py",
    "omni_speech/serve/controller.py",
    "omni_speech/serve/model_worker.py",
    "omni_speech/serve/gradio_web_server.py",
    "omni_speech/infer/__init__.py",
    "omni_speech/infer/inference.py",
    "omni_speech/infer/run.sh"
]

# Define required packages
required_packages = [
    "torch",
    "transformers",
    "gradio",
    "fastapi",
    "uvicorn",
    "pydantic",
    "numpy",
    "tqdm"
]

def check_directory_structure():
    """Check if all required directories exist."""
    print("Checking directory structure...")
    missing_dirs = []
    
    for dir_path in required_dirs:
        if not os.path.isdir(dir_path):
            missing_dirs.append(dir_path)
    
    if missing_dirs:
        print(f"❌ Missing directories: {', '.join(missing_dirs)}")
        return False
    else:
        print("✅ All required directories exist.")
        return True

def check_required_files():
    """Check if all required files exist."""
    print("Checking required files...")
    missing_files = []
    
    for file_path in required_files:
        if not os.path.isfile(file_path):
            missing_files.append(file_path)
    
    if missing_files:
        print(f"❌ Missing files: {', '.join(missing_files)}")
        return False
    else:
        print("✅ All required files exist.")
        return True

def check_packages():
    """Check if all required packages are installed."""
    print("Checking required packages...")
    missing_packages = []
    
    for package in required_packages:
        if importlib.util.find_spec(package) is None:
            missing_packages.append(package)
    
    if missing_packages:
        print(f"❌ Missing packages: {', '.join(missing_packages)}")
        return False
    else:
        print("✅ All required packages are installed.")
        return True

def check_python_version():
    """Check if Python version is compatible."""
    print("Checking Python version...")
    major, minor = sys.version_info[:2]
    
    if major != 3 or minor < 10:
        print(f"❌ Incompatible Python version: {major}.{minor}. Python 3.10+ is required.")
        return False
    else:
        print(f"✅ Python version is compatible: {major}.{minor}")
        return True

def main():
    """Run all checks."""
    print("🔍 Checking LLaMA-Omni setup...")
    print("-" * 50)
    
    checks = [
        check_directory_structure(),
        check_required_files(),
        check_packages(),
        check_python_version()
    ]
    
    print("-" * 50)
    
    if all(checks):
        print("✅ All checks passed! LLaMA-Omni is set up correctly.")
        print("🚀 Run 'python app.py' to start the application.")
    else:
        print("❌ Some checks failed. Please fix the issues before running the application.")

if __name__ == "__main__":
    main()