#!/usr/bin/env python3
"""
Test script for Textilindo AI Assistant deployment
Run this to verify your setup before deploying to Hugging Face Spaces
"""

import os
import sys
import json
import requests
from pathlib import Path

def test_file_structure():
    """Test if all required files exist"""
    print("🔍 Testing file structure...")
    
    required_files = [
        "app.py",
        "Dockerfile", 
        "requirements.txt",
        "configs/system_prompt.md",
        "configs/training_config.yaml",
        "templates/chat.html"
    ]
    
    required_dirs = [
        "data",
        "configs",
        "templates"
    ]
    
    missing_files = []
    missing_dirs = []
    
    for file in required_files:
        if not Path(file).exists():
            missing_files.append(file)
    
    for dir in required_dirs:
        if not Path(dir).exists():
            missing_dirs.append(dir)
    
    if missing_files:
        print(f"❌ Missing files: {missing_files}")
        return False
    
    if missing_dirs:
        print(f"❌ Missing directories: {missing_dirs}")
        return False
    
    print("✅ All required files and directories exist")
    return True

def test_data_files():
    """Test if data files exist and are valid"""
    print("🔍 Testing data files...")
    
    data_dir = Path("data")
    if not data_dir.exists():
        print("❌ Data directory not found")
        return False
    
    jsonl_files = list(data_dir.glob("*.jsonl"))
    if not jsonl_files:
        print("❌ No JSONL files found in data directory")
        return False
    
    print(f"✅ Found {len(jsonl_files)} JSONL files:")
    for file in jsonl_files:
        print(f"  - {file.name}")
    
    # Test one JSONL file
    test_file = jsonl_files[0]
    try:
        with open(test_file, 'r', encoding='utf-8') as f:
            lines = f.readlines()
        
        if not lines:
            print(f"❌ {test_file.name} is empty")
            return False
        
        # Test first line
        first_line = lines[0].strip()
        if first_line:
            json.loads(first_line)
            print(f"✅ {test_file.name} contains valid JSON")
        
        print(f"✅ {test_file.name} has {len(lines)} lines")
        return True
        
    except json.JSONDecodeError as e:
        print(f"❌ {test_file.name} contains invalid JSON: {e}")
        return False
    except Exception as e:
        print(f"❌ Error reading {test_file.name}: {e}")
        return False

def test_config_files():
    """Test if configuration files are valid"""
    print("🔍 Testing configuration files...")
    
    # Test system prompt
    prompt_file = Path("configs/system_prompt.md")
    if not prompt_file.exists():
        print("❌ System prompt file not found")
        return False
    
    try:
        with open(prompt_file, 'r', encoding='utf-8') as f:
            content = f.read()
        
        if 'SYSTEM_PROMPT' not in content:
            print("⚠️  SYSTEM_PROMPT not found in system_prompt.md")
        else:
            print("✅ System prompt file is valid")
        
    except Exception as e:
        print(f"❌ Error reading system prompt: {e}")
        return False
    
    # Test training config
    config_file = Path("configs/training_config.yaml")
    if not config_file.exists():
        print("❌ Training config file not found")
        return False
    
    try:
        import yaml
        with open(config_file, 'r') as f:
            config = yaml.safe_load(f)
        
        required_fields = ['model_name', 'model_path', 'dataset_path']
        for field in required_fields:
            if field not in config:
                print(f"❌ Missing field in config: {field}")
                return False
        
        print("✅ Training configuration is valid")
        return True
        
    except Exception as e:
        print(f"❌ Error reading training config: {e}")
        return False

def test_app_import():
    """Test if the app can be imported"""
    print("🔍 Testing app import...")
    
    try:
        # Add current directory to path
        sys.path.insert(0, '.')
        
        # Try to import the app
        import app
        print("✅ App module imported successfully")
        
        # Test if FastAPI app exists
        if hasattr(app, 'app'):
            print("✅ FastAPI app found")
        else:
            print("❌ FastAPI app not found")
            return False
        
        return True
        
    except ImportError as e:
        print(f"❌ Error importing app: {e}")
        return False
    except Exception as e:
        print(f"❌ Unexpected error: {e}")
        return False

def test_environment():
    """Test environment variables"""
    print("🔍 Testing environment...")
    
    # Check if HUGGINGFACE_API_KEY is set
    api_key = os.getenv('HUGGINGFACE_API_KEY')
    if api_key:
        print("✅ HUGGINGFACE_API_KEY is set")
    else:
        print("⚠️  HUGGINGFACE_API_KEY not set (will use mock responses)")
    
    # Check Python version
    python_version = sys.version_info
    if python_version >= (3, 8):
        print(f"✅ Python version {python_version.major}.{python_version.minor} is compatible")
    else:
        print(f"❌ Python version {python_version.major}.{python_version.minor} is too old (need 3.8+)")
        return False
    
    return True

def test_dependencies():
    """Test if required dependencies can be imported"""
    print("🔍 Testing dependencies...")
    
    required_modules = [
        'fastapi',
        'uvicorn',
        'pydantic',
        'requests',
        'huggingface_hub'
    ]
    
    missing_modules = []
    
    for module in required_modules:
        try:
            __import__(module)
            print(f"✅ {module}")
        except ImportError:
            missing_modules.append(module)
            print(f"❌ {module}")
    
    if missing_modules:
        print(f"❌ Missing modules: {missing_modules}")
        print("Install with: pip install " + " ".join(missing_modules))
        return False
    
    return True

def main():
    """Run all tests"""
    print("🧪 Textilindo AI Assistant - Deployment Test")
    print("=" * 50)
    
    tests = [
        test_file_structure,
        test_data_files,
        test_config_files,
        test_environment,
        test_dependencies,
        test_app_import
    ]
    
    passed = 0
    total = len(tests)
    
    for test in tests:
        try:
            if test():
                passed += 1
            print()
        except Exception as e:
            print(f"❌ Test failed with error: {e}")
            print()
    
    print("=" * 50)
    print(f"📊 Test Results: {passed}/{total} tests passed")
    
    if passed == total:
        print("🎉 All tests passed! Your setup is ready for deployment.")
        print("\n📋 Next steps:")
        print("1. Run: ./deploy_to_hf_space.sh")
        print("2. Or manually deploy to Hugging Face Spaces")
        print("3. Set environment variables in your space settings")
        print("4. Test your deployed application")
    else:
        print("❌ Some tests failed. Please fix the issues above before deploying.")
        return 1
    
    return 0

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