import os
import subprocess
import logging
from ultralytics import YOLO
from config import CONFIG

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='{"time": "%(asctime)s", "level": "%(levelname)s", "message": "%(message)s"}'
)
logger = logging.getLogger(__name__)

def check_ffmpeg():
    """Verify FFmpeg is installed."""
    try:
        subprocess.run(["ffmpeg", "-version"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
        logger.info("FFmpeg is available.")
        return True
    except (subprocess.CalledProcessError, FileNotFoundError):
        logger.error("FFmpeg is not installed or not found in PATH.")
        return False

def check_model_files():
    """Verify model files exist or can be downloaded."""
    model_path = CONFIG["MODEL_PATH"]
    fallback_model = CONFIG["FALLBACK_MODEL"]
    
    if not os.path.isfile(model_path):
        logger.warning(f"Custom model {model_path} not found. Checking fallback.")
        if not os.path.isfile(fallback_model):
            logger.info(f"Downloading fallback model: {fallback_model}")
            try:
                import torch
                torch.hub.download_url_to_file(
                    'https://github.com/ultralytics/assets/releases/download/v8.3.0/yolov8n.pt',
                    fallback_model
                )
                logger.info(f"Downloaded {fallback_model}")
            except Exception as e:
                logger.error(f"Failed to download {fallback_model}: {e}")
                return False
    else:
        logger.info(f"Custom model found: {model_path}")
    return True

def test_model_inference():
    """Test YOLO model inference with a dummy image."""
    try:
        model_path = CONFIG["MODEL_PATH"] if os.path.isfile(CONFIG["MODEL_PATH"]) else CONFIG["FALLBACK_MODEL"]
        model = YOLO(model_path)
        # Create a dummy image (1x1 pixel)
        import numpy as np
        dummy_image = np.zeros((1, 1, 3), dtype=np.uint8)
        results = model(dummy_image, conf=0.1, verbose=False)
        logger.info("Model inference test passed.")
        return True
    except Exception as e:
        logger.error(f"Model inference test failed: {e}")
        return False

def run_integration_tests():
    """Run all integration tests."""
    logger.info("Starting integration tests...")
    
    tests = [
        ("FFmpeg availability", check_ffmpeg),
        ("Model files check", check_model_files),
        ("Model inference test", test_model_inference)
    ]
    
    all_passed = True
    for test_name, test_func in tests:
        logger.info(f"Running test: {test_name}")
        result = test_func()
        if not result:
            logger.error(f"Test failed: {test_name}")
            all_passed = False
        else:
            logger.info(f"Test passed: {test_name}")
    
    return all_passed

if __name__ == "__main__":
    success = run_integration_tests()
    exit(0 if success else 1)