#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MIGAN 인페인팅 모듈 테스트 스크립트

이 스크립트는 MIGAN 모듈의 기본 사용법과 설정을 보여줍니다.
실제 운영 환경에서는 celery_worker.py를 통해 사용하세요.
"""

import os
import sys
import numpy as np
import cv2
from pathlib import Path

# 현재 디렉토리를 Python 경로에 추가
current_dir = Path(__file__).parent
sys.path.insert(0, str(current_dir))

try:
    from migan_inpainting_module import MIGANInpaintingModule, build_migan_from_toggle, quick_migan_inpaint
    from loggerModule import Logger
except ImportError as e:
    print(f"모듈 import 실패: {e}")
    print("worker 디렉토리에서 실행하거나 PYTHONPATH를 설정하세요.")
    sys.exit(1)


def create_test_image_and_mask():
    """테스트용 이미지와 마스크 생성"""
    # 테스트 이미지 생성 (800x600, 컬러)
    image = np.random.randint(0, 255, (600, 800, 3), dtype=np.uint8)
    
    # 중앙에 텍스트 영역 시뮬레이션 (여러 개의 사각형)
    mask = np.zeros((600, 800), dtype=np.uint8)
    
    # 텍스트 영역들 (마스크에서 255는 텍스트 영역)
    rectangles = [
        (100, 100, 200, 50),   # x, y, w, h
        (300, 200, 150, 40),
        (500, 300, 180, 60),
        (200, 400, 220, 45),
    ]
    
    for x, y, w, h in rectangles:
        mask[y:y+h, x:x+w] = 255
    
    return image, mask


def test_basic_usage():
    """기본 사용법 테스트"""
    print("=== MIGAN 인페인팅 모듈 기본 테스트 ===")
    
    # 로거 생성
    logger = Logger()
    
    # 테스트 이미지 및 마스크 생성
    print("1. 테스트 이미지 및 마스크 생성...")
    image, mask = create_test_image_and_mask()
    print(f"   이미지 크기: {image.shape}")
    print(f"   마스크 크기: {mask.shape}")
    print(f"   마스크 영역: {np.sum(mask > 0)} 픽셀")
    
    # MIGAN 모듈 초기화
    print("\n2. MIGAN 모듈 초기화...")
    try:
        # 기본 설정으로 모듈 생성
        config = {
            'onnx_path': '/app/worker/models/migan_pipeline_v2.onnx',
            'use_cuda': False,  # CPU 모드로 테스트
            'use_tensorrt': False,
            'max_image_size': 1024,
            'output_max_width': 600,
        }
        
        inpainter = MIGANInpaintingModule(logger=logger, config=config)
        print("   MIGAN 모듈 초기화 성공")
        
        # 처리 통계 확인
        print("\n3. 처리 통계 분석...")
        stats = inpainter.get_processing_stats(image, mask)
        for key, value in stats.items():
            print(f"   {key}: {value}")
        
        # 인페인팅 실행 (모델 파일이 없으면 실패할 것임)
        print("\n4. MIGAN 인페인팅 실행...")
        try:
            result = inpainter.inpaint_with_migan(image, mask)
            print(f"   인페인팅 성공! 결과 이미지 크기: {result.shape}")
        except FileNotFoundError:
            print("   ⚠️ 모델 파일을 찾을 수 없습니다: /app/worker/models/migan_pipeline_v2.onnx")
            print("   모델 파일을 올바른 위치에 배치한 후 다시 시도하세요.")
        except Exception as e:
            print(f"   인페인팅 실패: {e}")
        
        # 메모리 정리
        inpainter.cleanup_memory()
        
    except Exception as e:
        print(f"   MIGAN 모듈 초기화 실패: {e}")


def test_toggle_states_usage():
    """toggle_states를 사용한 설정 테스트"""
    print("\n=== toggle_states 설정 테스트 ===")
    
    # 로거 생성
    logger = Logger()
    
    # toggle_states 설정 예시들
    test_configs = [
        {
            "name": "CPU 모드",
            "toggle_states": {
                "inpaint_method": "migan",
                "migan_use_cuda": False,
                "migan_use_tensorrt": False,
                "migan_max_image_size": 800,
                "migan_output_max_width": 600,
            }
        },
        {
            "name": "GPU 모드 (TensorRT 비활성화)",
            "toggle_states": {
                "inpaint_method": "migan",
                "migan_use_cuda": True,
                "migan_use_tensorrt": False,
                "migan_max_image_size": 1600,
                "migan_output_max_width": 800,
            }
        },
        {
            "name": "최고 성능 모드 (TensorRT + FP16)",
            "toggle_states": {
                "inpaint_method": "migan",
                "migan_use_cuda": True,
                "migan_use_tensorrt": True,
                "migan_trt_fp16_enable": True,
                "migan_trt_engine_cache_enable": True,
                "migan_max_image_size": 2048,
                "migan_output_max_width": 1200,
            }
        }
    ]
    
    for i, config in enumerate(test_configs, 1):
        print(f"\n{i}. {config['name']} 테스트")
        print(f"   설정: {config['toggle_states']}")
        
        try:
            # toggle_states로부터 MIGAN 모듈 생성
            inpainter = build_migan_from_toggle(config['toggle_states'], logger)
            print("   모듈 생성 성공")
            
            # 설정 확인
            print(f"   CUDA 사용: {inpainter.default_config['use_cuda']}")
            print(f"   TensorRT 사용: {inpainter.default_config['use_tensorrt']}")
            print(f"   최대 이미지 크기: {inpainter.default_config['max_image_size']}")
            
        except Exception as e:
            print(f"   모듈 생성 실패: {e}")


def test_quick_inpaint():
    """빠른 인페인팅 함수 테스트"""
    print("\n=== 빠른 인페인팅 함수 테스트 ===")
    
    # 테스트 이미지 생성
    image, mask = create_test_image_and_mask()
    
    # 빠른 인페인팅 설정
    config = {
        'use_cuda': False,
        'use_tensorrt': False,
        'max_image_size': 512,
    }
    
    print("1. quick_migan_inpaint 함수 실행...")
    try:
        result = quick_migan_inpaint(image, mask, config=config)
        print(f"   성공! 결과 크기: {result.shape}")
    except FileNotFoundError:
        print("   ⚠️ 모델 파일을 찾을 수 없습니다.")
    except Exception as e:
        print(f"   실패: {e}")


def test_error_handling():
    """에러 처리 테스트"""
    print("\n=== 에러 처리 테스트 ===")
    
    logger = Logger()
    
    # 잘못된 모델 경로로 테스트
    print("1. 잘못된 모델 경로 테스트...")
    config = {
        'onnx_path': '/nonexistent/path/model.onnx',
        'use_cuda': False,
    }
    
    try:
        inpainter = MIGANInpaintingModule(logger=logger, config=config)
        image, mask = create_test_image_and_mask()
        result = inpainter.inpaint_with_migan(image, mask)
        print("   예상과 다르게 성공함")
    except Exception as e:
        print(f"   예상된 에러 발생: {type(e).__name__}")
    
    # 잘못된 입력 데이터 테스트
    print("\n2. 잘못된 입력 데이터 테스트...")
    try:
        config = {'use_cuda': False, 'use_tensorrt': False}
        inpainter = MIGANInpaintingModule(logger=logger, config=config)
        
        # 잘못된 형태의 이미지/마스크
        bad_image = np.zeros((100, 100), dtype=np.uint8)  # 2D 이미지
        bad_mask = np.zeros((50, 50), dtype=np.uint8)     # 다른 크기
        
        result = inpainter.inpaint_with_migan(bad_image, bad_mask)
        print("   예상과 다르게 성공함")
    except Exception as e:
        print(f"   예상된 에러 발생: {type(e).__name__}")


def show_usage_examples():
    """사용 예시 출력"""
    print("\n=== 사용 예시 ===")
    
    print("\n1. Celery worker에서 ROI vs MIGAN 선택:")
    print("""
# ROI 기반 인페인팅 (기본)
toggle_states = {
    "inpaint_method": "roi"
}

# MIGAN 인페인팅
toggle_states = {
    "inpaint_method": "migan",
    "migan_use_cuda": True,
    "migan_trt_fp16_enable": True
}
""")
    
    print("\n2. 직접 MIGAN 모듈 사용:")
    print("""
from worker.migan_inpainting_module import MIGANInpaintingModule

# 모듈 생성
inpainter = MIGANInpaintingModule(config={
    'use_cuda': True,
    'max_image_size': 1600
})

# 인페인팅 실행
result = inpainter.inpaint_with_migan(image, mask)
""")
    
    print("\n3. toggle_states로 모듈 생성:")
    print("""
from worker.migan_inpainting_module import build_migan_from_toggle

toggle_states = {
    "migan_use_tensorrt": True,
    "migan_trt_fp16_enable": True,
    "migan_max_image_size": 2048
}

inpainter = build_migan_from_toggle(toggle_states)
""")


def main():
    """메인 테스트 실행"""
    print("MIGAN 인페인팅 모듈 테스트 시작")
    print("=" * 50)
    
    # 환경 정보 출력
    print(f"Python 버전: {sys.version}")
    print(f"작업 디렉토리: {os.getcwd()}")
    print(f"NumPy 버전: {np.__version__}")
    print(f"OpenCV 버전: {cv2.__version__}")
    
    # 모델 파일 존재 확인
    model_path = "/app/worker/models/migan_pipeline_v2.onnx"
    if os.path.exists(model_path):
        print(f"✅ 모델 파일 발견: {model_path}")
    else:
        print(f"⚠️ 모델 파일 없음: {model_path}")
        print("   모델 파일을 배치한 후 테스트하면 실제 인페인팅이 실행됩니다.")
    
    # 테스트 실행
    try:
        test_basic_usage()
        test_toggle_states_usage()
        test_quick_inpaint()
        test_error_handling()
        show_usage_examples()
        
    except KeyboardInterrupt:
        print("\n사용자에 의해 중단됨")
    except Exception as e:
        print(f"\n예상치 못한 에러: {e}")
        import traceback
        traceback.print_exc()
    
    print("\n" + "=" * 50)
    print("테스트 완료")


if __name__ == "__main__":
    main()