#!/usr/bin/env python3
"""
Testes básicos para ImgBoost.
Execute: python test_imgboost.py
"""

import numpy as np
import cv2
from imgboost import Engine
from imgboost.processors import ImageFilters, DocumentProcessor, SuperResolution
from imgboost.utils import ImageIO


def create_test_image(size=(640, 480), noise_level=20):
    """Cria uma imagem de teste com ruído."""
    # Criar imagem base
    img = np.random.randint(100, 200, (size[1], size[0], 3), dtype=np.uint8)
    
    # Adicionar algumas formas
    cv2.rectangle(img, (50, 50), (200, 200), (255, 255, 255), -1)
    cv2.circle(img, (400, 300), 80, (0, 0, 0), -1)
    cv2.putText(img, "Test Image", (150, 400), cv2.FONT_HERSHEY_SIMPLEX, 
                2, (255, 0, 0), 3)
    
    # Adicionar ruído
    noise = np.random.normal(0, noise_level, img.shape).astype(np.int16)
    img = np.clip(img.astype(np.int16) + noise, 0, 255).astype(np.uint8)
    
    return img


def test_image_io():
    """Testar carregamento e salvamento de imagens."""
    print("🧪 Teste 1: ImageIO")
    
    io = ImageIO()
    
    # Criar imagem de teste
    test_img = create_test_image()
    
    # Salvar
    io.save(test_img, 'test_output.jpg', quality=95)
    print("  ✅ Salvamento OK")
    
    # Carregar
    loaded = io.load('test_output.jpg')
    assert loaded is not None, "Falha ao carregar imagem"
    print("  ✅ Carregamento OK")
    
    # Validar
    assert io.validate_image('test_output.jpg'), "Validação falhou"
    print("  ✅ Validação OK")
    
    print("✅ ImageIO: PASSOU\n")


def test_filters():
    """Testar filtros de imagem."""
    print("🧪 Teste 2: Filtros")
    
    filters = ImageFilters()
    img = create_test_image()
    
    # Denoise
    denoised = filters.denoise(img)
    assert denoised.shape == img.shape, "Shape incorreto após denoise"
    print("  ✅ Denoise OK")
    
    # Contrast
    enhanced = filters.enhance_contrast(img)
    assert enhanced.shape == img.shape, "Shape incorreto após enhance"
    print("  ✅ Contrast OK")
    
    # Sharpen
    sharpened = filters.sharpen(img)
    assert sharpened.shape == img.shape, "Shape incorreto após sharpen"
    print("  ✅ Sharpen OK")
    
    # Dark UI
    dark_enhanced = filters.enhance_dark_ui(img)
    assert dark_enhanced.shape == img.shape, "Shape incorreto após dark UI"
    print("  ✅ Dark UI OK")
    
    print("✅ Filtros: PASSOU\n")


def test_document_processor():
    """Testar processador de documentos."""
    print("🧪 Teste 3: Processador de Documentos")
    
    doc = DocumentProcessor()
    img = create_test_image()
    
    # Enhance text
    enhanced = doc.enhance_text(img)
    assert enhanced is not None, "Falha no enhance_text"
    print("  ✅ Enhance text OK")
    
    # Binarize
    binary = doc.binarize(img, method='adaptive')
    assert binary is not None, "Falha na binarização"
    assert len(binary.shape) == 2, "Binarização deve retornar grayscale"
    print("  ✅ Binarização OK")
    
    # OCR enhancement
    ocr_ready = doc.enhance_for_ocr(img)
    assert ocr_ready is not None, "Falha no enhance_for_ocr"
    print("  ✅ OCR enhancement OK")
    
    print("✅ Document Processor: PASSOU\n")


def test_super_resolution():
    """Testar super-resolução."""
    print("🧪 Teste 4: Super-resolução")
    
    sr = SuperResolution()
    img = create_test_image(size=(320, 240))
    
    # Upscale 2x
    upscaled = sr.upscale(img, scale=2)
    expected_shape = (480, 640, 3)
    assert upscaled.shape == expected_shape, f"Shape incorreto: {upscaled.shape} != {expected_shape}"
    print("  ✅ Upscale 2x OK")
    
    # Upscale 4x
    upscaled_4x = sr.upscale(img, scale=4)
    expected_shape_4x = (960, 1280, 3)
    assert upscaled_4x.shape == expected_shape_4x, f"Shape incorreto: {upscaled_4x.shape} != {expected_shape_4x}"
    print("  ✅ Upscale 4x OK")
    
    print("✅ Super-resolução: PASSOU\n")


def test_engine():
    """Testar engine principal."""
    print("🧪 Teste 5: Engine")
    
    # Criar imagem de teste
    io = ImageIO()
    test_img = create_test_image()
    io.save(test_img, 'test_input.jpg')
    
    engine = Engine(output_dir='test_output')
    
    # Testar modo general
    output = engine.run('test_input.jpg', mode='general', upscale=False)
    assert output.exists(), "Arquivo de saída não foi criado"
    print("  ✅ Modo general OK")
    
    # Testar modo text
    output = engine.run('test_input.jpg', mode='text', upscale=False)
    assert output.exists(), "Arquivo de saída não foi criado"
    print("  ✅ Modo text OK")
    
    # Testar modo dark
    output = engine.run('test_input.jpg', mode='dark', upscale=False)
    assert output.exists(), "Arquivo de saída não foi criado"
    print("  ✅ Modo dark OK")
    
    # Testar com upscale
    output = engine.run('test_input.jpg', mode='general', upscale=True)
    assert output.exists(), "Arquivo de saída não foi criado"
    print("  ✅ Upscale OK")
    
    print("✅ Engine: PASSOU\n")


def run_all_tests():
    """Executar todos os testes."""
    print("=" * 60)
    print("🧪 Executando testes do ImgBoost")
    print("=" * 60)
    print()
    
    try:
        test_image_io()
        test_filters()
        test_document_processor()
        test_super_resolution()
        test_engine()
        
        print("=" * 60)
        print("✅ TODOS OS TESTES PASSARAM!")
        print("=" * 60)
        return True
        
    except Exception as e:
        print("=" * 60)
        print(f"❌ TESTE FALHOU: {e}")
        print("=" * 60)
        import traceback
        traceback.print_exc()
        return False


if __name__ == '__main__':
    success = run_all_tests()
    exit(0 if success else 1)
