#!/usr/bin/env python3
"""
LLaMA-Omni2 Debug Script
-------------------
This script helps diagnose issues with LLaMA-Omni2 installation and imports.
It checks:
1. Python environment
2. Module locations
3. Import capabilities
4. Script locations
"""

import os
import sys
import importlib
import subprocess

def print_section(title):
    """Print a section header for better readability"""
    print("\n" + "=" * 50)
    print(f" {title} ".center(50, "="))
    print("=" * 50)

def find_module_in_paths(module_name, paths=None):
    """Find all occurrences of a module in the specified paths"""
    if paths is None:
        paths = sys.path
    
    found_locations = []
    for path in paths:
        potential_path = os.path.join(path, module_name)
        if os.path.exists(potential_path):
            found_locations.append(potential_path)
    
    return found_locations

def find_scripts(script_name, search_dirs=None):
    """Find scripts by name in the specified directories"""
    if search_dirs is None:
        search_dirs = [
            '/code',
            '/tmp/LLaMA-Omni2',
            '/usr/local/lib/python3.10/site-packages',
            '/home/user'
        ]
    
    found_scripts = []
    
    for search_dir in search_dirs:
        if not os.path.exists(search_dir):
            continue
            
        for root, dirs, files in os.walk(search_dir):
            # Skip .git and other large dirs
            dirs[:] = [d for d in dirs if d not in ('.git', 'node_modules')]
            
            if script_name in files:
                found_scripts.append(os.path.join(root, script_name))
    
    return found_scripts

def check_pip_installed():
    """Check if llama_omni2 is properly installed via pip"""
    try:
        result = subprocess.run(['pip', 'show', 'llama-omni2'], 
                                 capture_output=True, text=True)
        if result.returncode == 0:
            print("LLaMA-Omni2 is installed via pip:")
            print(result.stdout)
        else:
            print("LLaMA-Omni2 is NOT installed via pip")
    except Exception as e:
        print(f"Error checking pip installation: {e}")

def main():
    # 1. Environment Information
    print_section("ENVIRONMENT INFORMATION")
    print(f"Python Executable: {sys.executable}")
    print(f"Python Version: {sys.version}")
    print(f"Working Directory: {os.getcwd()}")
    
    # 2. PYTHONPATH
    print_section("PYTHONPATH")
    pythonpath = os.environ.get('PYTHONPATH', 'Not set')
    print(f"PYTHONPATH Environment Variable: {pythonpath}")
    
    # 3. sys.path
    print_section("sys.path")
    for i, path in enumerate(sys.path):
        print(f"{i}: {path}")
    
    # 4. Check if llama_omni2 is pip-installed
    print_section("PIP INSTALLATION")
    check_pip_installed()
    
    # 5. Find llama_omni2 in sys.path
    print_section("LLAMA_OMNI2 MODULE LOCATIONS")
    found_locations = find_module_in_paths('llama_omni2')
    if found_locations:
        print("Found llama_omni2 module in the following locations:")
        for loc in found_locations:
            print(f"  - {loc}")
    else:
        print("Could not find llama_omni2 module in sys.path!")
    
    # 6. Try to import llama_omni2
    print_section("IMPORT TEST")
    try:
        import llama_omni2
        print(f"Successfully imported llama_omni2 from: {llama_omni2.__file__}")
        
        # Check if key modules exist
        modules_to_check = [
            'llama_omni2.serve.controller',
            'llama_omni2.serve.model_worker',
            'llama_omni2.serve.gradio_web_server'
        ]
        
        for module in modules_to_check:
            try:
                importlib.import_module(module)
                print(f"✅ Successfully imported {module}")
            except ImportError as e:
                print(f"❌ Failed to import {module}: {e}")
                
    except ImportError as e:
        print(f"Failed to import llama_omni2: {e}")
    
    # 7. Find core script files
    print_section("SCRIPT LOCATIONS")
    scripts_to_find = ['controller.py', 'model_worker.py', 'gradio_web_server.py']
    
    for script in scripts_to_find:
        found_scripts = find_scripts(script)
        if found_scripts:
            print(f"Found {script} at:")
            for path in found_scripts:
                print(f"  - {path}")
        else:
            print(f"Could not find {script}")
    
    # 8. Test running the scripts directly
    print_section("DIRECT SCRIPT EXECUTION TEST")
    
    for script in scripts_to_find:
        found_scripts = find_scripts(script)
        if found_scripts:
            script_path = found_scripts[0]
            print(f"Testing if {script_path} can be executed...")
            try:
                # Just import the script module directly to see if it loads
                result = subprocess.run(
                    [sys.executable, '-c', f"import importlib.util; spec = importlib.util.spec_from_file_location('test', '{script_path}'); module = importlib.util.module_from_spec(spec); spec.loader.exec_module(module); print('Successfully loaded {script}')"],
                    capture_output=True, text=True, timeout=5
                )
                if result.returncode == 0:
                    print(f"✅ Script can be imported: {script_path}")
                    print(result.stdout)
                else:
                    print(f"❌ Script import failed: {script_path}")
                    print(result.stderr)
            except subprocess.TimeoutExpired:
                print(f"⚠️ Script import timed out: {script_path}")
            except Exception as e:
                print(f"❌ Error testing script: {e}")

if __name__ == "__main__":
    main()