# 이미지 업로드 및 처리 섹션#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
한국어 OCR 텍스트 추출기
Google Gemini AI를 활용한 고정밀 한국어 문자 인식 애플리케이션
"""

import gradio as gr
import base64
import requests
import json
from PIL import Image
import io
import os
from typing import Optional, Tuple
import re
import time
import random

# 모듈 import 확인
try:
    import gradio as gr
    print("✅ Gradio 모듈 정상 로드됨")
except ImportError as e:
    print(f"❌ Gradio 모듈 로드 실패: {e}")
    print("pip install gradio==4.44.0 명령어로 설치해주세요")
    exit(1)

try:
    from PIL import Image
    print("✅ Pillow 모듈 정상 로드됨")
except ImportError as e:
    print(f"❌ Pillow 모듈 로드 실패: {e}")
    print("pip install Pillow==10.4.0 명령어로 설치해주세요")
    exit(1)

try:
    import requests
    print("✅ Requests 모듈 정상 로드됨")
except ImportError as e:
    print(f"❌ Requests 모듈 로드 실패: {e}")
    print("pip install requests==2.32.3 명령어로 설치해주세요")
    exit(1)

class KoreanOCRApp:
    def __init__(self):
        self.api_key = None
        self.project_id = None
        
    def set_credentials(self, api_key: str, project_id: str) -> str:
        """API 키와 프로젝트 ID 설정 및 검증"""
        if not api_key or not project_id:
            return "❌ API 키와 프로젝트 ID를 모두 입력해주세요."
        
        # 프로젝트 ID 검증 (영문, 숫자, 하이픈만 허용)
        if not re.match(r'^[a-z0-9\-]+$', project_id.strip()):
            return "❌ 유효하지 않은 프로젝트 ID 형식입니다. 영문 소문자, 숫자, 하이픈만 사용 가능합니다."
        
        self.api_key = api_key.strip()
        self.project_id = project_id.strip()
        return "✅ 인증 정보가 설정되었습니다."
    
    def optimize_image_for_api(self, image: Image.Image) -> Image.Image:
        """API 호출을 위한 이미지 최적화"""
        # 이미지 크기 최적화 (토큰 사용량 감소)
        max_dimension = 1024  # 더 작은 크기로 제한
        
        # 현재 이미지 크기 확인
        width, height = image.size
        
        # 큰 이미지일 경우 리사이즈
        if width > max_dimension or height > max_dimension:
            image.thumbnail((max_dimension, max_dimension), Image.Resampling.LANCZOS)
        
        # RGBA를 RGB로 변환 (파일 크기 감소)
        if image.mode == 'RGBA':
            background = Image.new('RGB', image.size, (255, 255, 255))
            background.paste(image, mask=image.split()[-1])
            image = background
        elif image.mode != 'RGB':
            image = image.convert('RGB')
        
        return image
    def encode_image_to_base64(self, image: Image.Image) -> str:
        """이미지를 base64로 인코딩 (최적화된 버전)"""
        # 이미지 최적화
        image = self.optimize_image_for_api(image)
        
        buffer = io.BytesIO()
        # JPEG 형식으로 저장하여 파일 크기 최적화 (품질 80으로 낮춤)
        image.save(buffer, format='JPEG', quality=80, optimize=True)
        image_bytes = buffer.getvalue()
        
        # 파일 크기 확인
        size_mb = len(image_bytes) / (1024 * 1024)
        if size_mb > 3:  # 3MB 초과 시 추가 최적화
            buffer = io.BytesIO()
            image.save(buffer, format='JPEG', quality=60, optimize=True)
            image_bytes = buffer.getvalue()
        
        return base64.b64encode(image_bytes).decode('utf-8')
    
    def call_gemini_api_with_retry(self, image_base64: str, max_retries: int = 3, initial_delay: float = 2.0) -> str:
        """재시도 로직이 포함된 Gemini API 호출"""
        if not self.api_key:
            return "❌ 먼저 API 키를 설정해주세요."
        
        # Google AI Studio API 엔드포인트 사용
        url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-pro:generateContent?key={self.api_key}"
        
        headers = {
            "Content-Type": "application/json"
        }
        
        payload = {
            "contents": [{
                "parts": [
                    {
                        "text": """이 이미지에 포함된 모든 한국어 텍스트를 정확하게 추출해주세요. 
                        다음 규칙을 따라주세요:
                        1. 이미지에서 발견되는 모든 한국어 텍스트를 순서대로 추출
                        2. 텍스트의 위치나 레이아웃을 최대한 보존
                        3. 줄바꿈과 문단 구분을 명확히 표시
                        4. 특수문자, 숫자, 영어가 포함되어 있다면 그대로 유지
                        5. 읽기 어려운 부분이 있다면 [불분명] 표시
                        
                        추출된 텍스트만 반환해주세요."""
                    },
                    {
                        "inline_data": {
                            "mime_type": "image/jpeg",
                            "data": image_base64
                        }
                    }
                ]
            }],
            "generationConfig": {
                "temperature": 0.1,
                "topP": 0.8,
                "topK": 40,
                "maxOutputTokens": 8192
            },
            "safetySettings": [
                {
                    "category": "HARM_CATEGORY_HARASSMENT",
                    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
                },
                {
                    "category": "HARM_CATEGORY_HATE_SPEECH", 
                    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
                },
                {
                    "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
                    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
                },
                {
                    "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
                    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
                }
            ]
        }
        
        for attempt in range(max_retries):
            try:
                response = requests.post(url, headers=headers, json=payload, timeout=60)
                
                if response.status_code == 401:
                    return "❌ API 키가 유효하지 않습니다. Google AI Studio에서 발급받은 올바른 API 키를 입력해주세요."
                elif response.status_code == 403:
                    return "❌ API 접근 권한이 없습니다. Gemini API가 활성화되어 있는지 확인해주세요."
                elif response.status_code == 429:
                    # 429 에러 시 재시도 로직
                    if attempt < max_retries - 1:
                        delay = initial_delay * (2 ** attempt) + random.uniform(0.5, 1.5)  # 지수 백오프 + 랜덤 지터
                        return f"⏳ API 호출 한도를 초과했습니다. {delay:.1f}초 후 자동으로 재시도합니다... (시도 {attempt + 1}/{max_retries})"
                    else:
                        return """❌ API 호출 한도를 초과했습니다. 
                        
📌 해결 방법:
1. 잠시 기다린 후 다시 시도 (1-2분 권장)
2. Google AI Studio에서 할당량 확인
3. 유료 계정으로 업그레이드 고려
4. 이미지 크기를 줄여서 재시도

💡 팁: 높은 해상도의 이미지는 더 많은 토큰을 사용합니다."""
                
                response.raise_for_status()
                
                result = response.json()
                
                if "candidates" in result and len(result["candidates"]) > 0:
                    content = result["candidates"][0]["content"]["parts"][0]["text"]
                    return content.strip()
                elif "error" in result:
                    error_msg = result['error'].get('message', '알 수 없는 오류')
                    if "quota" in error_msg.lower() or "limit" in error_msg.lower():
                        if attempt < max_retries - 1:
                            delay = initial_delay * (2 ** attempt) + random.uniform(0.5, 1.5)
                            time.sleep(delay)
                            continue
                    return f"❌ API 오류: {error_msg}"
                else:
                    return "❌ 텍스트를 추출할 수 없습니다. 이미지에 한국어 텍스트가 포함되어 있는지 확인해주세요."
                    
            except requests.exceptions.RequestException as e:
                if "429" in str(e) and attempt < max_retries - 1:
                    delay = initial_delay * (2 ** attempt) + random.uniform(0.5, 1.5)
                    time.sleep(delay)
                    continue
                return f"❌ API 호출 오류: {str(e)}"
            except json.JSONDecodeError:
                return "❌ API 응답 파싱 오류가 발생했습니다."
            except KeyError as e:
                return f"❌ 예상치 못한 API 응답 형식: {str(e)}"
            except Exception as e:
                return f"❌ 알 수 없는 오류: {str(e)}"
        
        return "❌ 최대 재시도 횟수를 초과했습니다. 잠시 후 다시 시도해주세요."
    
    def call_vertex_ai_api(self, image_base64: str) -> str:
        """Vertex AI API 호출 (서비스 계정 키 사용)"""
        if not self.api_key or not self.project_id:
            return "❌ 먼저 API 키와 프로젝트 ID를 설정해주세요."
        
        location = "us-central1"
        url = f"https://{location}-aiplatform.googleapis.com/v1/projects/{self.project_id}/locations/{location}/publishers/google/models/gemini-1.5-pro:generateContent"
        
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        payload = {
            "contents": [{
                "role": "user",
                "parts": [
                    {
                        "text": """이 이미지에 포함된 모든 한국어 텍스트를 정확하게 추출해주세요. 
                        다음 규칙을 따라주세요:
                        1. 이미지에서 발견되는 모든 한국어 텍스트를 순서대로 추출
                        2. 텍스트의 위치나 레이아웃을 최대한 보존
                        3. 줄바꿈과 문단 구분을 명확히 표시
                        4. 특수문자, 숫자, 영어가 포함되어 있다면 그대로 유지
                        5. 읽기 어려운 부분이 있다면 [불분명] 표시
                        
                        추출된 텍스트만 반환해주세요."""
                    },
                    {
                        "inline_data": {
                            "mime_type": "image/jpeg",
                            "data": image_base64
                        }
                    }
                ]
            }],
            "generation_config": {
                "temperature": 0.1,
                "top_p": 0.8,
                "top_k": 40,
                "max_output_tokens": 8192
            }
        }
        
        try:
            response = requests.post(url, headers=headers, json=payload, timeout=60)
            
            if response.status_code == 401:
                return "❌ 인증 오류: Access Token이 유효하지 않거나 만료되었습니다."
            elif response.status_code == 403:
                return "❌ 권한 오류: Vertex AI API 접근 권한이 없습니다."
            elif response.status_code == 404:
                return "❌ 프로젝트 ID가 올바르지 않거나 Vertex AI API가 활성화되지 않았습니다."
            
            response.raise_for_status()
            
            result = response.json()
            
            if "candidates" in result and len(result["candidates"]) > 0:
                content = result["candidates"][0]["content"]["parts"][0]["text"]
                return content.strip()
            else:
                return "❌ 텍스트를 추출할 수 없습니다. 이미지에 한국어 텍스트가 포함되어 있는지 확인해주세요."
                
        except requests.exceptions.RequestException as e:
            return f"❌ API 호출 오류: {str(e)}"
        except Exception as e:
            return f"❌ 알 수 없는 오류: {str(e)}"
    
    def process_image(self, image: Optional[Image.Image], api_key: str, project_id: str, api_type: str) -> Tuple[Optional[Image.Image], str]:
        """이미지 처리 및 OCR 수행"""
        if image is None:
            return None, "❌ 이미지를 업로드해주세요."
        
        # 인증 정보 설정
        if api_type == "Google AI Studio":
            if not api_key:
                return image, "❌ Google AI Studio API 키를 입력해주세요."
            self.api_key = api_key.strip()
        else:  # Vertex AI
            auth_result = self.set_credentials(api_key, project_id)
            if "❌" in auth_result:
                return image, auth_result
        
        try:
            # 이미지 최적화 (토큰 사용량 감소를 위해)
            image = self.optimize_image_for_api(image)
            
            # 이미지를 base64로 인코딩
            image_base64 = self.encode_image_to_base64(image)
            
            # API 타입에 따라 호출
            if api_type == "Google AI Studio":
                extracted_text = self.call_gemini_api_with_retry(image_base64)
            else:
                extracted_text = self.call_vertex_ai_api(image_base64)
            
            # 결과 반환
            return image, extracted_text
            
        except Exception as e:
            return image, f"❌ 이미지 처리 중 오류가 발생했습니다: {str(e)}"

# 전역 앱 인스턴스
ocr_app = KoreanOCRApp()

def create_interface():
    """Gradio 인터페이스 생성"""
    
    # CSS 스타일링
    css = """
    .gradio-container {
        font-family: 'Noto Sans KR', sans-serif;
    }
    
    .main-header {
        text-align: center;
        color: #2c3e50;
        margin-bottom: 20px;
    }
    
    .info-box {
        background-color: #e8f4fd;
        border: 1px solid #bee5eb;
        border-radius: 8px;
        padding: 15px;
        margin: 10px 0;
    }
    
    .warning-box {
        background-color: #fff3cd;
        border: 1px solid #ffeaa7;
        border-radius: 8px;
        padding: 15px;
        margin: 10px 0;
        color: #856404;
    }
    """
    
    with gr.Blocks(css=css, title="한국어 OCR - Gemini AI") as interface:
        
        gr.Markdown("""
        # 🔍 한국어 OCR 텍스트 추출기
        ### Google Gemini AI를 활용한 고정밀 한국어 문자 인식
        
        이미지에서 한국어 텍스트를 정확하게 추출합니다. 문서, 간판, 손글씨 등 다양한 형태의 한국어를 인식할 수 있습니다.
        """, elem_classes="main-header")
        
        # API 선택
        gr.Markdown("## 🔧 API 설정")
        
        api_type = gr.Radio(
            choices=["Google AI Studio", "Vertex AI"],
            value="Google AI Studio",
            label="사용할 API 선택",
            info="Google AI Studio는 개인 사용자용, Vertex AI는 기업용"
        )
        
        # 인증 정보 입력 섹션
        with gr.Row():
            with gr.Column(scale=2):
                api_key_input = gr.Textbox(
                    label="API 키 / Access Token",
                    placeholder="Google AI Studio API 키 또는 Vertex AI Access Token",
                    type="password",
                    lines=1
                )
            with gr.Column(scale=1):
                project_id_input = gr.Textbox(
                    label="프로젝트 ID (Vertex AI만)",
                    placeholder="Google Cloud 프로젝트 ID",
                    lines=1
                )
        
        # API 설정 가이드
        with gr.Accordion("📖 API 설정 가이드 및 할당량 정보", open=False):
            gr.Markdown("""
            ### Google AI Studio API (권장)
            1. [Google AI Studio](https://aistudio.google.com/)에 접속
            2. "Get API Key" 클릭
            3. API 키 생성 및 복사
            4. 위의 "API 키" 필드에 붙여넣기
            
            **📊 무료 할당량 (Google AI Studio):**
            - 분당 15회 요청
            - 일일 1,500회 요청
            - 분당 100만 토큰
            - 일일 5천만 토큰
            
            ### Vertex AI API (고급 사용자용)
            1. [Google Cloud Console](https://console.cloud.google.com/)에서 프로젝트 생성
            2. Vertex AI API 활성화
            3. 서비스 계정 생성 및 키 다운로드
            4. `gcloud auth application-default login` 또는 Access Token 발급
            5. API 키와 프로젝트 ID 입력
            
            ### ⚠️ 할당량 초과 시 해결 방법
            1. **잠시 대기**: 1-2분 후 다시 시도
            2. **이미지 최적화**: 더 작은 크기의 이미지 사용
            3. **사용량 분산**: 여러 번 나누어서 처리
            4. **유료 계정**: Google Cloud 유료 계정으로 업그레이드
            
            ### 💡 토큰 절약 팁
            - 이미지 해상도: 1024x1024 이하 권장
            - 파일 형식: JPEG 사용 (PNG보다 작음)
            - 불필요한 배경 제거
            - 텍스트 영역만 크롭하여 업로드
            """, elem_classes="warning-box")
        
        # 할당량 상태 표시
        with gr.Row():
            gr.Markdown("""
            ### 📊 현재 상태
            
            **무료 할당량 (Google AI Studio):**
            - ⏱️ 분당 15회 요청 제한
            - 📅 일일 1,500회 요청 제한
            - 🔢 고해상도 이미지는 더 많은 토큰 사용
            
            **💡 할당량 절약 팁:**
            - 이미지 크기를 1024x1024 이하로 유지
            - 텍스트가 있는 부분만 크롭하여 업로드
            - 연속적인 요청 간 1-2초 간격 유지
            """, elem_classes="info-box")
        gr.Markdown("## 📤 이미지 업로드 및 텍스트 추출")
        
        with gr.Row():
            with gr.Column(scale=1):
                input_image = gr.Image(
                    label="📁 이미지 업로드",
                    type="pil",
                    sources=["upload", "clipboard"],
                    interactive=True
                )
                
                process_btn = gr.Button(
                    "🔍 텍스트 추출 시작", 
                    variant="primary",
                    size="lg"
                )
                
            with gr.Column(scale=1):
                output_image = gr.Image(
                    label="📋 업로드된 이미지 확인",
                    type="pil",
                    interactive=False
                )
        
        # 추출된 텍스트 출력
        gr.Markdown("## 📝 추출된 텍스트")
        
        extracted_text = gr.Textbox(
            label="인식된 한국어 텍스트",
            placeholder="추출된 텍스트가 여기에 표시됩니다...",
            lines=10,
            max_lines=20,
            interactive=True,
            show_copy_button=True
        )
        
        # 이벤트 핸들러
        process_btn.click(
            fn=ocr_app.process_image,
            inputs=[input_image, api_key_input, project_id_input, api_type],
            outputs=[output_image, extracted_text],
            show_progress=True
        )
        
        # 사용 팁
        gr.Markdown("""
        ### 💡 사용 팁
        
        **📸 이미지 품질:**
        - 선명하고 해상도가 높은 이미지 사용
        - 충분한 조명과 대비
        - 텍스트가 수평으로 배치된 이미지 권장
        
        **📄 지원 형식:**
        - **이미지 형식:** PNG, JPEG, WebP
        - **최대 크기:** 4MB (자동 리사이즈)
        - **인식 언어:** 한국어, 영어, 숫자, 특수문자
        
        **🔒 보안:**
        - API 키는 세션 동안만 임시 저장
        - 이미지는 서버에 저장되지 않음
        - 개인정보가 포함된 이미지 사용 시 주의
        
        **⚡ 성능:**
        - Google AI Studio: 빠르고 안정적 (권장)
        - Vertex AI: 기업용 고급 기능
        """)
    
    return interface

# 메인 실행
if __name__ == "__main__":
    try:
        print("🚀 한국어 OCR 애플리케이션을 시작합니다...")
        
        # 인터페이스 생성
        demo = create_interface()
        
        print("✅ 인터페이스 생성 완료")
        print("🌐 서버를 시작합니다...")
        
        # 서버 실행
        demo.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=True,
            debug=True,
            show_error=True,
            inbrowser=True
        )
        
    except Exception as e:
        print(f"❌ 애플리케이션 시작 중 오류 발생: {e}")
        print("\n🔧 문제 해결 방법:")
        print("1. pip install gradio==4.44.0 Pillow==10.4.0 requests==2.32.3")
        print("2. Python 버전 확인 (3.8 이상 필요)")
        print("3. 가상환경 사용 권장")
        print("4. 네트워크 연결 상태 확인")
        raise