#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Request_AI_Server의 _use_backup_rembg 함수 테스트 코드
"""

import os
import sys
import cv2
import logging
import numpy as np
from datetime import datetime

# 현재 파일의 디렉토리를 기준으로 src 폴더를 Python 경로에 추가
current_dir = os.path.dirname(os.path.abspath(__file__))
src_dir = os.path.dirname(current_dir)  # src 폴더
project_dir = os.path.dirname(src_dir)  # 프로젝트 루트
sys.path.insert(0, src_dir)

try:
    from modules.request_inpaint import Request_AI_Server
    from modules.gpu_utils import GPUManager
except ImportError as e:
    print(f"모듈 import 실패: {e}")
    print(f"현재 디렉토리: {current_dir}")
    print(f"src 디렉토리: {src_dir}")
    print(f"sys.path: {sys.path[:3]}...")
    sys.exit(1)

class SimpleLogger:
    """간단한 로거 클래스"""
    
    def __init__(self, name="TestLogger"):
        self.name = name
        
    def log(self, message, level=logging.INFO, exc_info=False):
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        level_name = logging.getLevelName(level)
        print(f"[{timestamp}] [{level_name}] {message}")
        if exc_info:
            import traceback
            traceback.print_exc()

def test_backup_rembg():
    """_use_backup_rembg 함수 테스트"""
    
    print("=== Request_AI_Server _use_backup_rembg 함수 테스트 ===")
    
    # 1. 테스트 이미지 경로 확인
    test_image_path = os.path.join(current_dir, "1.jpg")
    if not os.path.exists(test_image_path):
        print(f"❌ 테스트 이미지가 없습니다: {test_image_path}")
        print("1.jpg 파일을 다음 경로에 복사해주세요:")
        print(f"   {current_dir}")
        return False
    
    print(f"✅ 테스트 이미지 발견: {test_image_path}")
    
    # 2. 이미지 로드
    try:
        image_data = cv2.imread(test_image_path)
        if image_data is None:
            print(f"❌ 이미지 로드 실패: {test_image_path}")
            return False
        
        height, width = image_data.shape[:2]
        print(f"✅ 이미지 로드 성공: {width}x{height}")
        
    except Exception as e:
        print(f"❌ 이미지 로드 중 오류: {e}")
        return False
    
    # 3. 로거 및 GPU 관리자 초기화
    logger = SimpleLogger("BackupRemBGTest")
    
    try:
        gpu_manager = GPUManager(logger=logger)
        # DirectML GPU 모드로 강제 설정 (테스트용)
        gpu_manager.can_use_cuda = True
        print(f"✅ GPU 관리자 초기화 성공: CUDA 사용 가능 = {gpu_manager.can_use_cuda} (DirectML 강제 활성화)")
    except Exception as e:
        print(f"⚠️ GPU 관리자 초기화 실패, None으로 설정: {e}")
        gpu_manager = None
    
    # 4. Request_AI_Server 초기화
    try:
        # 로컬 rembg 모델 경로 설정
        local_rembg_model_path = os.path.join(src_dir, "modules", "rembg_models", "birefnet-general-lite.onnx")
        if not os.path.exists(local_rembg_model_path):
            print(f"⚠️ 로컬 rembg 모델이 없습니다: {local_rembg_model_path}")
            local_rembg_model_path = None
        else:
            print(f"✅ 로컬 rembg 모델 발견: {local_rembg_model_path}")
        
        request_ai_server = Request_AI_Server(
            logger=logger,
            inpaint_server_url="http://localhost:8008",  # 테스트용 (실제로는 사용하지 않음)
            rembg_server_url=None,
            gpu_manager=gpu_manager,
            local_rembg_model_path=local_rembg_model_path
        )
        print("✅ Request_AI_Server 초기화 성공")
        
    except Exception as e:
        print(f"❌ Request_AI_Server 초기화 실패: {e}")
        return False
    
    # 5. _use_backup_rembg 함수 테스트
    print("\n--- 백업 rembg 모듈 테스트 시작 ---")
    
    test_cases = [
        {"model_name": "birefnet-general-lite", "object_ratio": 0.75, "force_cpu": False, "desc": "DirectML GPU 모드 (문제)"},
        {"model_name": "birefnet-general-lite", "object_ratio": 0.75, "force_cpu": True, "desc": "CPU 모드 (정상)"},
        {"model_name": "birefnet-general-lite", "object_ratio": 1, "force_cpu": True, "desc": "CPU 모드 + 다른 비율"},
    ]
    
    results = []
    for i, test_case in enumerate(test_cases):
        print(f"\n테스트 케이스 {i+1}: {test_case}")
        
        try:
            start_time = datetime.now()
            
            result_img = request_ai_server._use_backup_rembg(
                image_data=image_data,
                model_name=test_case["model_name"],
                object_ratio=test_case["object_ratio"],
                debug_save=True,
                debug_prefix=f"test_{i+1}_{test_case['desc'].replace(' ', '_')}",
                force_cpu=test_case["force_cpu"]
            )
            
            end_time = datetime.now()
            processing_time = (end_time - start_time).total_seconds()
            
            if result_img is not None:
                # 결과 이미지 저장
                # 안전한 파일명 생성 (한글 및 특수문자 제거)
                safe_desc = test_case['desc'].replace('DirectML GPU 모드 (문제)', 'DirectML_GPU_problem').replace('CPU 모드 (정상)', 'CPU_normal').replace('CPU 모드 + 다른 비율', 'CPU_different_ratio')
                output_filename = f"backup_rembg_result_{i+1}_{safe_desc}_ratio{int(test_case['object_ratio']*100)}.png"
                output_path = os.path.join(current_dir, output_filename)
                
                success = cv2.imwrite(output_path, result_img)
                if success:
                    result_height, result_width = result_img.shape[:2]
                    print(f"  ✅ 성공! 결과 이미지 크기: {result_width}x{result_height}")
                    print(f"  📁 저장 경로: {output_path}")
                    print(f"  ⏱️ 처리 시간: {processing_time:.2f}초")
                    
                    results.append({
                        "test_case": i+1,
                        "success": True,
                        "output_path": output_path,
                        "size": (result_width, result_height),
                        "processing_time": processing_time
                    })
                else:
                    print(f"  ❌ 결과 이미지 저장 실패: {output_path}")
                    results.append({"test_case": i+1, "success": False, "error": "이미지 저장 실패"})
            else:
                print(f"  ❌ 함수 실행 실패: 결과 이미지가 None")
                results.append({"test_case": i+1, "success": False, "error": "결과 이미지 None"})
                
        except Exception as e:
            print(f"  ❌ 테스트 케이스 {i+1} 실행 중 오류: {e}")
            results.append({"test_case": i+1, "success": False, "error": str(e)})
    
    # 6. 테스트 결과 요약
    print("\n=== 테스트 결과 요약 ===")
    success_count = sum(1 for r in results if r.get("success", False))
    total_count = len(results)
    
    print(f"총 테스트: {total_count}개")
    print(f"성공: {success_count}개")
    print(f"실패: {total_count - success_count}개")
    
    for result in results:
        if result.get("success", False):
            print(f"  ✅ 케이스 {result['test_case']}: {result['size'][0]}x{result['size'][1]}, {result['processing_time']:.2f}초")
        else:
            print(f"  ❌ 케이스 {result['test_case']}: {result.get('error', '알 수 없는 오류')}")
    
    return success_count > 0

if __name__ == "__main__":
    success = test_backup_rembg()
    
    if success:
        print("\n🎉 테스트 완료! 결과 이미지를 확인해보세요.")
        print(f"📁 결과 파일 위치: {current_dir}")
    else:
        print("\n💥 테스트 실패!")
        
    sys.exit(0 if success else 1)