#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
이미지 번역 프로젝트 메인 진입점
중국어 상품 이미지를 한국어로 번역하는 시스템
CPU 전용 모드로 동작합니다.
"""

import os
import sys
import argparse
from pathlib import Path

# CPU만 사용하도록 환경 변수 설정
os.environ['CUDA_VISIBLE_DEVICES'] = ''
os.environ['OMP_NUM_THREADS'] = '4'  # OpenMP 스레드 수 제한

from modules.ocr_module import OCRModule
# from modules.ocr_m2 import OCRModule2
from modules.mask_module import MaskModule
from modules.translation_module import TranslationModule
from modules.inpainting_module import InpaintingModule
from modules.text_rendering_module import TextRenderingModule
from modules.output_module import OutputModule
from gpt import GPTClient

class ImageTranslator:
    def __init__(self):
        """이미지 번역기 초기화 (CPU 전용)"""
        print("🖥️  CPU 전용 모드로 이미지 번역기 초기화 중...")
        self.ocr = OCRModule()
        # self.ocr2 = OCRModule2()
        self.mask_processor = MaskModule()
        self.translator = TranslationModule()
        self.inpainter = InpaintingModule()
        self.text_renderer = TextRenderingModule()
        self.output_manager = OutputModule()
        self.gpt_client = GPTClient()

        print("✅ 이미지 번역기 초기화 완료")
        
    def translate_image(self, image_path, output_dir="output", inpainting_method="cv2", ocr_method="polygon"):
        """
        이미지 번역 파이프라인 실행 (결과 이미지 포함)
        
        Args:
            image_path (str): 입력 이미지 경로
            output_dir (str): 출력 디렉토리
            inpainting_method (str): 인페인팅 방법 (cv2_telea, cv2_ns, lama, stable_diffusion, edge_inpaint, patchmatch, iopaint)
            ocr_method (str): OCR 감지 방법 (polygon, bbox, expanded_bbox, rotated_bbox, contour)
        """
        print(f"🚀 이미지 번역 시작: {image_path}")
        print(f"📋 설정 - OCR: {ocr_method}, 인페인팅: {inpainting_method}")
        print(f"🖥️  실행 모드: CPU 전용")
        
        # 출력 디렉토리 생성
        os.makedirs(output_dir, exist_ok=True)
        
        # 1. OCR로 텍스트 감지 (시각화 이미지 포함)
        print("\n1️⃣ OCR 텍스트 감지 중...")
        ocr_results, ocr_visualization_path = self.ocr.detect_text_with_visualization(
            image_path, method=ocr_method, output_dir=output_dir
        )
        
        if not ocr_results:
            print("❌ 텍스트가 감지되지 않았습니다.")
            return
            
        # 2. 마스크 생성 (마스크 이미지 포함)
        print("\n2️⃣ 마스크 생성 중...")
        mask_option = "basic"
        # mask_option = "morphology"
        # mask_option = "adaptive"
        # mask_option = "contour"
        # mask_option = "ultra_blur"
        # mask_option = "feathered"
        

        masks, mask_image_path = self.mask_processor.create_masks_with_save(
            image_path, ocr_results, output_dir, mask_option=mask_option
        )
        
        # 3. 텍스트 번역
        print("\n3️⃣ 텍스트 번역 중...")
        translated_texts = self.gpt_translate_texts(ocr_results, self.gpt_client)
        for original, translated in zip([r['text'] for r in ocr_results], translated_texts):
            print(f"   '{original}' → '{translated}'")
        
        print(f"\n4️⃣ 인페인팅 중... (방법: {inpainting_method})")
        inpainted_image = self.inpainter.inpaint(image_path, masks, method=inpainting_method)
        
        
        # 5. 텍스트 렌더링
        print("\n5️⃣ 번역된 텍스트 렌더링 중...")
        final_image = self.text_renderer.render_text(
            inpainted_image, ocr_results, translated_texts, font_path="HakgyoansimDunggeunmisoTTFB.ttf"
            # inpainted_image, ocr_results, [tr.text for tr in translated_texts], font_path="HakgyoansimDunggeunmisoTTFB.ttf"

        )
        
        # 6. 결과 출력
        print("\n6️⃣ 결과 생성 중...")
        self.output_manager.create_comparison_output(
            image_path, inpainted_image, final_image, 
            output_dir, inpainting_method
        )
        
        # 7. 결과 요약
        print(f"\n✅ 번역 완료! 결과 파일들:")
        print(f"   📁 출력 폴더: {output_dir}")
        if ocr_visualization_path:
            print(f"   🔍 OCR 감지 이미지: {os.path.basename(ocr_visualization_path)}")
        if mask_image_path:
            print(f"   🎭 마스크 이미지: {os.path.basename(mask_image_path)}")
        print(f"   🖼️  최종 번역 결과: 비교 이미지 파일")
        print(f"   📊 메타데이터: JSON 파일")
        
        return {
            'ocr_visualization': ocr_visualization_path,
            'mask_image': mask_image_path,
            'output_dir': output_dir
        }

    def gpt_translate_texts(self, ocr_results, gpt_client):
        """
        OCR 결과에서 추출한 텍스트들을 GPT를 이용해 한 번에 번역
        Args:
            ocr_results (list): OCR 결과 리스트 (각 원소는 dict, 'text' 키 포함)
            gpt_client (GPTClient): GPTClient 인스턴스
        Returns:
            list: 번역된 텍스트 리스트 (ocr_results와 순서 동일)
        """
        # 1. 번역할 텍스트 추출
        texts = [result['text'] for result in ocr_results]
        if not texts:
            return []

        # 2. 프롬프트 구성 (문맥 보존, 리스트 형태로 요청)
        prompt = (
            "다음 중국어 문장들을 한국어로 자연스럽고 의미가 잘 전달되게 번역해줘. "
            "순서와 개수는 반드시 그대로 유지하고, 결과는 JSON 배열(리스트)로만 반환해. "
            "중국어 리스트:\n" +
            str(texts)
        )

        # 3. GPT 번역 요청
        response = gpt_client.ask(prompt)
        # response가 리스트 형태로 올 수도 있고, 딕셔너리로 올 수도 있음
        if isinstance(response, list):
            return response
        elif isinstance(response, dict) and 'result' in response:
            return response['result']
        else:
            print("GPT 번역 결과 파싱 실패, 원본 반환")
            return texts

def test_modules():
    """각 모듈들의 기능을 개별적으로 테스트"""
    print("=== 모듈 테스트 시작 ===")
    
    # OCR 모듈 테스트
    print("\n1. OCR 모듈 테스트")
    ocr = OCRModule()
    ocr.test_module()
    
    # 마스크 모듈 테스트
    print("\n2. 마스크 모듈 테스트")
    mask_processor = MaskModule()
    mask_processor.test_module()
    
    # 번역 모듈 테스트
    print("\n3. 번역 모듈 테스트")
    translator = TranslationModule()
    translator.test_module()
    
    # 인페인팅 모듈 테스트
    print("\n4. 인페인팅 모듈 테스트")
    inpainter = InpaintingModule()
    inpainter.test_module()
    
    # 텍스트 렌더링 모듈 테스트
    print("\n5. 텍스트 렌더링 모듈 테스트")
    text_renderer = TextRenderingModule()
    text_renderer.test_module()
    
    # 출력 모듈 테스트
    print("\n6. 출력 모듈 테스트")
    output_manager = OutputModule()
    output_manager.test_module()
    
    print("\n=== 모듈 테스트 완료 ===")

def test_ocr_methods():
    """다양한 OCR 감지 방식 테스트"""
    print("🔍 === OCR 감지 방식 테스트 ===")
    
    ocr = OCRModule()
    test_image = "test_output/test_original.jpg"
    
    if not os.path.exists(test_image):
        print(f"❌ 테스트 이미지를 찾을 수 없습니다: {test_image}")
        return
    
    methods = ['polygon', 'bbox', 'expanded_bbox', 'rotated_bbox', 'contour']
    
    for method in methods:
        print(f"\n🎯 {method.upper()} 방식 테스트")
        print("-" * 50)
        results = ocr.detect_text(test_image, method=method)
        print(f"✅ {method} 방식 완료: {len(results)}개 텍스트 감지\n")

def test_inpainting_methods():
    """다양한 인페인팅 방식 테스트"""
    print("🎨 === 인페인팅 방식 테스트 ===")
    
    translator = ImageTranslator()
    test_image = "test_output/test_original.jpg"
    
    if not os.path.exists(test_image):
        print(f"❌ 테스트 이미지를 찾을 수 없습니다: {test_image}")
        return
    
    methods = ['cv2_telea', 'cv2_ns', 'lama', 'edge_inpaint']
    
    for method in methods:
        print(f"\n🖼️ {method.upper()} 방식으로 번역 테스트")
        print("-" * 50)
        output_dir = f"test_output/inpainting_{method}"
        os.makedirs(output_dir, exist_ok=True)
        translator.translate_image(test_image, output_dir, method)
        print(f"✅ {method} 방식 완료\n")

def test_combined_methods():
    """OCR과 인페인팅 방식 조합 테스트"""
    print("🔄 === OCR + 인페인팅 조합 테스트 ===")
    
    translator = ImageTranslator()
    test_image = "test_output/test_original.jpg"
    
    if not os.path.exists(test_image):
        print(f"❌ 테스트 이미지를 찾을 수 없습니다: {test_image}")
        return
    
    ocr_methods = ['polygon', 'expanded_bbox']
    inpainting_methods = ['cv2_telea', 'edge_inpaint']
    
    for ocr_method in ocr_methods:
        for inpainting_method in inpainting_methods:
            print(f"\n🎯 조합 테스트: {ocr_method} + {inpainting_method}")
            print("-" * 60)
            output_dir = f"test_output/combined_{ocr_method}_{inpainting_method}"
            os.makedirs(output_dir, exist_ok=True)
            translator.translate_image(test_image, output_dir, inpainting_method, ocr_method)
            print(f"✅ {ocr_method} + {inpainting_method} 조합 완료\n")

def benchmark_performance():
    """성능 벤치마크 테스트"""
    print("⏱️ === 성능 벤치마크 테스트 ===")
    
    import time
    
    translator = ImageTranslator()
    test_image = "test_output/test_original.jpg"
    
    if not os.path.exists(test_image):
        print(f"❌ 테스트 이미지를 찾을 수 없습니다: {test_image}")
        return
    
    methods_to_test = [
        ('polygon', 'cv2_telea'),
        ('expanded_bbox', 'cv2_telea'),
        ('polygon', 'edge_inpaint'),
        ('expanded_bbox', 'edge_inpaint')
    ]
    
    results = []
    
    for ocr_method, inpainting_method in methods_to_test:
        print(f"\n📊 벤치마크: {ocr_method} + {inpainting_method}")
        
        start_time = time.time()
        output_dir = f"test_output/benchmark_{ocr_method}_{inpainting_method}"
        os.makedirs(output_dir, exist_ok=True)
        
        translator.translate_image(test_image, output_dir, inpainting_method, ocr_method)
        
        end_time = time.time()
        elapsed = end_time - start_time
        
        results.append((f"{ocr_method}+{inpainting_method}", elapsed))
        print(f"⏱️ 소요 시간: {elapsed:.2f}초")
    
    print(f"\n📈 === 벤치마크 결과 요약 ===")
    for method_combo, elapsed in sorted(results, key=lambda x: x[1]):
        print(f"🏆 {method_combo}: {elapsed:.2f}초")

def main():
    parser = argparse.ArgumentParser(description='🌏 이미지 번역 시스템 - 중국어를 한국어로 번역')
    
    # 기본 옵션들
    parser.add_argument('--image', '-i', type=str, help='번역할 이미지 파일 경로')
    parser.add_argument('--output', '-o', type=str, default='output', help='출력 디렉토리 (기본값: output)')
    
    # 방법 선택 옵션들
    parser.add_argument('--inpainting', '-p', type=str, default='cv2_telea', 
                       choices=['cv2_telea', 'cv2_ns', 'lama', 'stable_diffusion', 'edge_inpaint', 'patchmatch', 'iopaint'],
                       help='인페인팅 방법 선택 (기본값: cv2_telea)')
    parser.add_argument('--ocr', '-r', type=str, default='polygon',
                       choices=['polygon', 'bbox', 'expanded_bbox', 'rotated_bbox', 'contour'],
                       help='OCR 감지 방법 선택 (기본값: polygon)')
    
    # 테스트 옵션들
    parser.add_argument('--test', '-t', action='store_true', help='전체 모듈 테스트 실행')
    parser.add_argument('--test-ocr', action='store_true', help='OCR 감지 방식들 테스트')
    parser.add_argument('--test-inpainting', action='store_true', help='인페인팅 방식들 테스트')
    parser.add_argument('--test-combined', action='store_true', help='OCR + 인페인팅 조합 테스트')
    parser.add_argument('--benchmark', '-b', action='store_true', help='성능 벤치마크 테스트')
    
    # 도움말 옵션
    parser.add_argument('--list-methods', '-l', action='store_true', help='사용 가능한 방법들 목록 출력')
    
    args = parser.parse_args()
    
    # 사용 가능한 방법들 목록 출력
    if args.list_methods:
        print("🔍 === 사용 가능한 OCR 감지 방법들 ===")
        print("• polygon       - 폴리곤 기반 감지 (기본값)")
        print("• bbox          - 바운딩 박스 감지")
        print("• expanded_bbox - 확장된 바운딩 박스 감지")
        print("• rotated_bbox  - 회전된 바운딩 박스 감지")
        print("• contour       - 컨투어 기반 감지")
        
        print("\n🎨 === 사용 가능한 인페인팅 방법들 ===")
        print("• cv2_telea     - OpenCV Telea 알고리즘 (기본값)")
        print("• cv2_ns        - OpenCV Navier-Stokes 알고리즘")
        print("• lama          - LaMa 딥러닝 모델")
        print("• stable_diffusion - Stable Diffusion 모델")
        print("• edge_inpaint  - 엣지 기반 인페인팅")
        
        print("\n💡 === 사용 예시 ===")
        print("python main.py -i image.jpg -r expanded_bbox -p edge_inpaint")
        print("python main.py --test-ocr")
        print("python main.py --benchmark")
        return
    
    # 테스트 옵션들 처리
    if args.test:
        test_modules()
        return
        
    if args.test_ocr:
        test_ocr_methods()
        return
        
    if args.test_inpainting:
        test_inpainting_methods()
        return
        
    if args.test_combined:
        test_combined_methods()
        return
        
    if args.benchmark:
        benchmark_performance()
        return
    
    # 이미지 번역 실행
    if not args.image:
        print("❌ 이미지 파일 경로를 지정해주세요.")
        print("💡 사용법: python main.py -i 이미지파일.jpg")
        print("📚 도움말: python main.py --help")
        print("📋 방법 목록: python main.py --list-methods")
        return
        
    if not os.path.exists(args.image):
        print(f"❌ 이미지 파일을 찾을 수 없습니다: {args.image}")
        return
        
    # 출력 디렉토리 생성
    os.makedirs(args.output, exist_ok=True)
    
    # 이미지 번역 실행
    print(f"🎯 선택된 설정:")
    print(f"   📁 입력 이미지: {args.image}")
    print(f"   📁 출력 디렉토리: {args.output}")
    print(f"   🔍 OCR 방법: {args.ocr}")
    print(f"   🎨 인페인팅 방법: {args.inpainting}")
    print()
    
    translator = ImageTranslator()
    translator.translate_image(args.image, args.output, args.inpainting, args.ocr)

if __name__ == "__main__":
    main()