#!/usr/bin/env python3
"""
test_model_download.py

Test script to verify the model download functionality.
This script runs download_models.py with various options and checks
that the expected files are created.
"""

import os
import sys
import subprocess
import argparse
from pathlib import Path

# Add project root to path for imports
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

def run_command(cmd):
    """Run a command and return stdout, stderr, and return code"""
    process = subprocess.Popen(
        cmd, 
        stdout=subprocess.PIPE, 
        stderr=subprocess.PIPE,
        universal_newlines=True,
        shell=True
    )
    stdout, stderr = process.communicate()
    return stdout, stderr, process.returncode

def test_download_all():
    """Test downloading all models"""
    print("\n=== Testing downloading all models ===")
    
    # Create a temporary directory for test downloads
    tmp_dir = Path('/tmp/morphguard_test_download')
    os.makedirs(tmp_dir, exist_ok=True)
    
    # Run the download script
    cmd = f"python {os.path.join(os.path.dirname(os.path.dirname(__file__)), 'download_models.py')} --skip-verification"
    stdout, stderr, rc = run_command(cmd)
    
    # Print output
    print("Command output:")
    print(stdout)
    
    if rc != 0:
        print(f"Error: Command returned non-zero exit code {rc}")
        print(f"stderr: {stderr}")
        return False
    
    # Check if the models directory was created
    if not os.path.exists('models'):
        print("Error: 'models' directory not created")
        return False
    
    print("Success: All models downloaded")
    return True

def test_download_specific():
    """Test downloading a specific model"""
    print("\n=== Testing downloading a specific model ===")
    
    # Remove existing model if any
    model_path = 'models/detector/morph_detector.pth'
    if os.path.exists(model_path):
        try:
            os.remove(model_path)
            print(f"Removed existing model {model_path}")
        except:
            print(f"Warning: Could not remove existing model {model_path}")
    
    # Run the download script
    cmd = f"python {os.path.join(os.path.dirname(os.path.dirname(__file__)), 'download_models.py')} --model morph_detector --skip-verification"
    stdout, stderr, rc = run_command(cmd)
    
    # Print output
    print("Command output:")
    print(stdout)
    
    if rc != 0:
        print(f"Error: Command returned non-zero exit code {rc}")
        print(f"stderr: {stderr}")
        return False
    
    # Check if the model was downloaded
    if not os.path.exists(model_path):
        print(f"Error: Model file {model_path} not created")
        return False
    
    print(f"Success: Model {model_path} downloaded")
    return True

def test_psp_download():
    """Test downloading pSp models"""
    print("\n=== Testing downloading pSp models ===")
    
    # Run the download script
    cmd = f"python {os.path.join(os.path.dirname(os.path.dirname(__file__)), 'download_models.py')} --psp --skip-verification"
    stdout, stderr, rc = run_command(cmd)
    
    # Print output
    print("Command output:")
    print(stdout)
    
    if rc != 0:
        print(f"Error: Command returned non-zero exit code {rc}")
        print(f"stderr: {stderr}")
        return False
    
    # Check if the pSp directory and a model file exist
    psp_model_path = 'pSp/pretrained_models/psp_ffhq_toonify.pt'
    if not os.path.exists(psp_model_path):
        print(f"Error: pSp model file {psp_model_path} not created")
        return False
    
    print(f"Success: pSp models downloaded")
    return True

def main():
    parser = argparse.ArgumentParser(description='Test the model download script')
    parser.add_argument('--all', action='store_true', help='Run all tests')
    parser.add_argument('--test-all', action='store_true', help='Test downloading all models')
    parser.add_argument('--test-specific', action='store_true', help='Test downloading a specific model')
    parser.add_argument('--test-psp', action='store_true', help='Test downloading pSp models')
    
    args = parser.parse_args()
    
    # If no tests specified, run all tests
    if not (args.test_all or args.test_specific or args.test_psp):
        args.all = True
    
    # Run tests
    results = {}
    
    if args.all or args.test_all:
        results['download_all'] = test_download_all()
    
    if args.all or args.test_specific:
        results['download_specific'] = test_download_specific()
    
    if args.all or args.test_psp:
        results['download_psp'] = test_psp_download()
    
    # Print summary
    print("\n=== Test Results ===")
    all_passed = True
    for test, result in results.items():
        if result:
            print(f"{test}: PASSED")
        else:
            print(f"{test}: FAILED")
            all_passed = False
    
    if all_passed:
        print("\nAll tests passed!")
        return 0
    else:
        print("\nSome tests failed. See above for details.")
        return 1

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