#!/usr/bin/env python3
"""
간단한 iopaint 호환성 테스트 서버
기본적인 API 엔드포인트들을 테스트할 수 있습니다.
"""
import base64
import time
import logging
from datetime import datetime
from typing import List
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import cv2
import numpy as np
from PIL import Image
import io

# 로깅 설정
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# FastAPI 앱 생성
app = FastAPI(title="iopaint 호환성 테스트 서버", version="1.0.0")

# CORS 미들웨어 추가
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 시작 시간 기록
start_time = time.time()

# Pydantic 모델들
class InpaintRequest(BaseModel):
    image: str  # base64 인코딩된 이미지
    mask: str   # base64 인코딩된 마스크
    model_name: str = "simple-lama"
    prompt: str = ""
    negative_prompt: str = ""
    sd_seed: int = -1
    num_inference_steps: int = 20
    guidance_scale: float = 7.5
    strength: float = 1.0

class RemoveBGRequest(BaseModel):
    image: str  # base64 인코딩된 이미지
    model_name: str = "rembg"

class InpaintResponse(BaseModel):
    success: bool
    image: str = None  # base64 인코딩된 결과 이미지
    error: str = None
    processing_time: float = None
    seed: int = None

class RemoveBGResponse(BaseModel):
    success: bool
    image: str = None  # base64 인코딩된 결과 이미지
    mask: str = None   # base64 인코딩된 마스크
    error: str = None
    processing_time: float = None

class HealthResponse(BaseModel):
    status: str
    timestamp: str
    version: str
    uptime: float

class ModelInfo(BaseModel):
    name: str
    type: str
    description: str = None
    supported_formats: List[str] = []
    max_image_size: int = None

class ServerConfigResponse(BaseModel):
    models: List[ModelInfo]
    max_file_size: int
    supported_formats: List[str]
    device: str
    is_jetson: bool

# 유틸리티 함수들
def decode_base64_to_image(base64_string: str, gray: bool = False):
    """base64 문자열을 이미지로 디코딩 (iopaint 호환)"""
    try:
        # base64 디코딩
        image_data = base64.b64decode(base64_string)
        
        # PIL Image로 변환
        image = Image.open(io.BytesIO(image_data))
        
        # numpy 배열로 변환
        if gray:
            image = image.convert('L')
            image_array = np.array(image)
        else:
            image_array = np.array(image.convert('RGB'))
        
        return image_array
    except Exception as e:
        logger.error(f"이미지 디코딩 실패: {e}")
        raise HTTPException(status_code=400, detail=f"이미지 디코딩 실패: {str(e)}")

def encode_image_to_base64(image_array: np.ndarray, format: str = "PNG"):
    """이미지를 base64로 인코딩 (iopaint 호환)"""
    try:
        # PIL Image로 변환
        if len(image_array.shape) == 3 and image_array.shape[2] == 3:
            image = Image.fromarray(image_array, 'RGB')
        else:
            image = Image.fromarray(image_array)
        
        # 바이트로 변환
        buffer = io.BytesIO()
        image.save(buffer, format=format)
        image_bytes = buffer.getvalue()
        
        # base64로 인코딩
        base64_string = base64.b64encode(image_bytes).decode('utf-8')
        return base64_string
    except Exception as e:
        logger.error(f"이미지 인코딩 실패: {e}")
        raise HTTPException(status_code=500, detail=f"이미지 인코딩 실패: {str(e)}")

def simulate_inpainting(image: np.ndarray, mask: np.ndarray, prompt: str = ""):
    """인페인팅 시뮬레이션 (실제 구현에서는 실제 모델 사용)"""
    # 마스크 영역을 약간 수정하여 시뮬레이션
    result = image.copy()
    
    # 마스크 영역을 약간 밝게 만들기
    mask_bool = mask > 127
    result[mask_bool] = np.clip(result[mask_bool] * 1.2, 0, 255).astype(np.uint8)
    
    # 약간의 노이즈 추가
    noise = np.random.randint(-20, 20, result.shape, dtype=np.int16)
    result = np.clip(result.astype(np.int16) + noise, 0, 255).astype(np.uint8)
    
    return result

def simulate_background_removal(image: np.ndarray):
    """배경 제거 시뮬레이션 (실제 구현에서는 실제 모델 사용)"""
    # 중앙 영역을 전경으로, 가장자리를 배경으로 가정
    height, width = image.shape[:2]
    center_x, center_y = width // 2, height // 2
    
    # 타원형 마스크 생성
    y, x = np.ogrid[:height, :width]
    mask = ((x - center_x) ** 2 / (width * 0.3) ** 2 + 
            (y - center_y) ** 2 / (height * 0.4) ** 2) <= 1
    
    # 마스크를 0-255 범위로 변환
    mask_uint8 = (mask * 255).astype(np.uint8)
    
    # 결과 이미지 (마스크 영역만 유지)
    result = image.copy()
    result[~mask] = [255, 255, 255]  # 배경을 흰색으로
    
    return result, mask_uint8

# API 엔드포인트들
@app.get("/health", response_model=HealthResponse)
async def health_check():
    """서버 상태 확인 (iopaint 호환)"""
    uptime = time.time() - start_time
    return HealthResponse(
        status="healthy",
        timestamp=datetime.now().isoformat(),
        version="1.0.0",
        uptime=uptime
    )

@app.get("/api/v1/server-config", response_model=ServerConfigResponse)
async def get_server_config():
    """서버 설정 정보 반환 (iopaint 호환)"""
    models = [
        ModelInfo(
            name="simple-lama",
            type="inpainting",
            description="Simple LAMA 인페인팅 모델 (테스트용)",
            supported_formats=["png", "jpg", "jpeg"],
            max_image_size=2048
        ),
        ModelInfo(
            name="migan",
            type="inpainting",
            description="MIGAN 인페인팅 모델 (테스트용)",
            supported_formats=["png", "jpg", "jpeg"],
            max_image_size=2048
        ),
        ModelInfo(
            name="rembg",
            type="rembg",
            description="Rembg 배경 제거 모델 (테스트용)",
            supported_formats=["png", "jpg", "jpeg"],
            max_image_size=2048
        )
    ]
    
    return ServerConfigResponse(
        models=models,
        max_file_size=25,
        supported_formats=["png", "jpg", "jpeg"],
        device="cuda",
        is_jetson=True
    )

@app.post("/api/v1/inpaint", response_model=InpaintResponse)
async def inpaint_image(request: InpaintRequest):
    """인페인팅 API (iopaint 호환)"""
    start_time_process = time.time()
    
    try:
        # base64 이미지 디코딩
        image = decode_base64_to_image(request.image)
        mask = decode_base64_to_image(request.mask, gray=True)
        
        # 이미지와 마스크 크기 검증
        if image.shape[:2] != mask.shape[:2]:
            raise HTTPException(
                status_code=400,
                detail=f"이미지 크기({image.shape[:2]})와 마스크 크기({mask.shape[:2]})가 일치하지 않습니다."
            )
        
        # 인페인팅 시뮬레이션
        result = simulate_inpainting(image, mask, request.prompt)
        
        # 결과를 base64로 인코딩
        result_base64 = encode_image_to_base64(result)
        
        processing_time = time.time() - start_time_process
        
        logger.info(f"인페인팅 성공: {image.shape}, 처리시간: {processing_time:.2f}초")
        
        return InpaintResponse(
            success=True,
            image=result_base64,
            processing_time=processing_time,
            seed=request.sd_seed
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"인페인팅 처리 실패: {e}")
        processing_time = time.time() - start_time_process
        
        return InpaintResponse(
            success=False,
            error=f"인페인팅 처리 실패: {str(e)}",
            processing_time=processing_time
        )

@app.post("/api/v1/remove_bg", response_model=RemoveBGResponse)
async def remove_background(request: RemoveBGRequest):
    """배경 제거 API (iopaint 호환)"""
    start_time_process = time.time()
    
    try:
        # base64 이미지 디코딩
        image = decode_base64_to_image(request.image)
        
        # 배경 제거 시뮬레이션
        result_image, result_mask = simulate_background_removal(image)
        
        # 결과를 base64로 인코딩
        result_base64 = encode_image_to_base64(result_image)
        mask_base64 = encode_image_to_base64(result_mask)
        
        processing_time = time.time() - start_time_process
        
        logger.info(f"배경 제거 성공: {image.shape}, 처리시간: {processing_time:.2f}초")
        
        return RemoveBGResponse(
            success=True,
            image=result_base64,
            mask=mask_base64,
            processing_time=processing_time
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"배경 제거 처리 실패: {e}")
        processing_time = time.time() - start_time_process
        
        return RemoveBGResponse(
            success=False,
            error=f"배경 제거 처리 실패: {str(e)}",
            processing_time=processing_time
        )

@app.get("/api/v1/samplers")
async def get_samplers():
    """사용 가능한 샘플러 목록 반환 (iopaint 호환)"""
    return [
        "euler", "euler_a", "heun", "dpm_2", "dpm_2_a",
        "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_a",
        "dpmpp_sde", "dpmpp_2m", "ddim", "uni_pc", "uni_pc_bh2"
    ]

@app.get("/")
async def root():
    """루트 엔드포인트"""
    return {
        "message": "iopaint 호환성 테스트 서버",
        "version": "1.0.0",
        "docs": "/docs",
        "health": "/health"
    }

if __name__ == "__main__":
    import uvicorn
    print("🚀 iopaint 호환성 테스트 서버 시작 중...")
    print("   📖 API 문서: http://localhost:8002/docs")
    print("   🏥 헬스 체크: http://localhost:8002/health")
    print("   🧪 테스트: http://localhost:8002/api/v1/server-config")
    
    uvicorn.run(app, host="0.0.0.0", port=8002, log_level="info")