# -*- coding: utf-8 -*-
import logging
import os
import cv2
import numpy as np

from ocr_module import OCRModule
from mask_module import MaskModule
from inpaint_module import InpaintModule


class SimpleLogger:
    """간단한 콘솔 로거 (OCRModule/MaskModule/InpaintModule과 동일 인터페이스)"""
    def log(self, msg, level=logging.INFO, exc_info=False):
        prefix = "[INFO]"
        if level >= logging.ERROR:
            prefix = "[ERROR]"
        elif level >= logging.WARNING:
            prefix = "[WARN]"
        print(f"{prefix} {msg}")
        if exc_info:
            import traceback
            traceback.print_exc()


def main():
    # 1) 샘플 이미지
    image_path = os.path.join("img", "1.jpg")
    if not os.path.isfile(image_path):
        raise FileNotFoundError(f"샘플 이미지가 존재하지 않습니다: {image_path}")

    logger = SimpleLogger()
    base_dir = os.getcwd()

    # 2) OCR
    ocr_module = OCRModule(logger=logger, base_dir=base_dir)
    ocr_results = ocr_module.detect_text(image_path, method="polygon")
    chinese_results = ocr_module.filter_chinese_text(ocr_results)

    print("\n=== 인식된 중국어 텍스트 ===")
    for idx, res in enumerate(chinese_results, 1):
        print(f"{idx:02d}. text = {res['text']}, conf = {res['confidence']:.3f}")

    # 3) 마스크 생성
    mask_module = MaskModule(logger=logger, base_dir=base_dir)
    mask_np = mask_module.create_masks(
        image_path=image_path,
        ocr_results=chinese_results,
        expansion_size=0,   # 필요시 조정
        blur_size=0,
        mask_option="basic"
    )
    if mask_np is None:
        logger.log("마스크 생성 실패, 종료합니다.", level=logging.ERROR)
        return

    # 저장 (필요시)
    mask_path = os.path.join("img", "1_mask.png")
    cv2.imwrite(mask_path, mask_np)
    logger.log(f"마스크 저장 완료: {mask_path}", level=logging.INFO)

    # 4) FastDeploy 인페인팅
    inpaint_module = InpaintModule(
        logger=logger,
        base_dir=base_dir,
        gpu_id=0,                      # GPU 번호
        backend="fd_sd"                # "fd_sd" or "opencv"
    )

    # 프롬프트를 비워두면 기본 "remove text..." 사용
    result_bgr = inpaint_module.inpaint(
        image_path=image_path,
        mask=mask_np,
        prompt="clean background, remove text, no letters",
        negative_prompt="text, letters, watermark, artifacts",
        steps=30,
        guidance_scale=7.5,
        strength=1.0,
        seed=42,
        cv_method="telea"
    )

    if result_bgr is None:
        logger.log("인페인팅 실패", level=logging.ERROR)
        return

    out_path = os.path.join("img", "1_inpainted.jpg")
    cv2.imwrite(out_path, result_bgr)
    logger.log(f"인페인팅 결과 저장: {out_path}", level=logging.INFO)

    # 5) 시각화(원본 + 마스크 + 결과) 간단 저장
    try:
        src = cv2.imread(image_path)
        mask_vis = cv2.cvtColor(mask_np, cv2.COLOR_GRAY2BGR)
        h, w = src.shape[:2]
        canvas = np.zeros((h, w * 3, 3), dtype=np.uint8)
        canvas[:, 0:w] = src
        canvas[:, w:2*w] = mask_vis
        canvas[:, 2*w:3*w] = result_bgr
        vis_path = os.path.join("img", "1_all_vis.jpg")
        cv2.imwrite(vis_path, canvas)
        logger.log(f"원본/마스크/결과 합성이미지 저장: {vis_path}", level=logging.INFO)
    except Exception as e:
        logger.log(f"시각화 실패: {e}", level=logging.WARNING)


if __name__ == "__main__":
    main()