main.py

# main.py - SiLIX Document Intelligence API (Chandra OCR 2 버전)
# 실행: docker compose up -d

from fastapi import FastAPI, UploadFile, File, HTTPException, BackgroundTasks
from fastapi.responses import JSONResponse
from pathlib import Path
import shutil
import uuid
import logging
from tool_chandra import async_process_document
from config_chandra import *
from typing import Dict
import time
import aiofiles


logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s | %(levelname)s | %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

app = FastAPI(
    title="SiLIX Document Intelligence API",
    description="문서 복호화 → 레이아웃 분석 → 마크다운/JSON 추출 엔진 (Chandra OCR 2)",
    version="2.0.0"
)

BASE_DIR = Path(__file__).parent
UPLOAD_DIR = BASE_DIR / "uploads"
UPLOAD_DIR.mkdir(exist_ok=True)


def cleanup_path(path: str):
    p = Path(path)
    try:
        if p.is_file():
            p.unlink()
        elif p.is_dir():
            shutil.rmtree(p, ignore_errors=True)
    except Exception as e:
        logger.warning(f"정리 실패: {path}, 오류: {e}")


@app.post("/process-file/")
async def process_file(
    file: UploadFile = File(...),
    background_tasks: BackgroundTasks = None,
    highqual: bool = False,
    use_large_model: bool = False   # 기존 인터페이스 호환 (Chandra에서는 무시)
):
    """
    업로드된 파일을 복호화하고, Chandra OCR 2로 레이아웃 분석을 수행하여
    마크다운, JSON, 시각화 이미지 등을 포함한 결과를 반환합니다.

    Args:
        file: 업로드된 파일 (PDF, Office, 이미지 등)
        background_tasks: 비동기 정리 작업
        highqual: 고해상도 처리 여부
        use_large_model: (호환용, Chandra에서는 무시됨)
    """
    content_type = file.content_type
    filename = file.filename
    job_id = str(uuid.uuid4())

    # --- 1. 파일 형식 및 크기 검증 ---
    if content_type not in ALLOWED_TYPES:
        raise HTTPException(status_code=400, detail=f"지원하지 않는 형식: {content_type}")

    content = await file.read()
    if len(content) > MAX_FILE_SIZE:
        raise HTTPException(status_code=413, detail="파일 크기 초과 (10MB 제한)")

    upload_path = UPLOAD_DIR / f"{job_id}_{filename}"

    try:
        # --- 2. 파일 저장 (비동기) ---
        async with aiofiles.open(upload_path, "wb") as f:
            await f.write(content)
        logger.info(f"[{job_id}] 업로드 완료: {filename}")

        # --- 3. Chandra OCR 2로 문서 처리 ---
        result = await async_process_document(
            file_path=str(upload_path),
            job_id=job_id,
            highqual=highqual,
            use_large_model=use_large_model
        )

        if result["status"] == "error":
            logger.error(f"[{job_id}] 문서 처리 실패: {result['message']}")
            raise HTTPException(status_code=500, detail=result["message"])

        # temp_dir 정리 등록
        temp_dir = result.get("temp_dir")
        if temp_dir and background_tasks:
            background_tasks.add_task(cleanup_path, temp_dir)

        # 응답 생성 (기존 인터페이스 완전 호환)
        response = {
            "job_id": job_id,
            "filename": filename,
            "status": "processed",
            "message": result["message"],
            "original_path": str(upload_path),
            "decrypted_path": result["decrypted_path"],
            "decrypted_filename": result["decrypted_filename"],
            "temp_dir": None,
            "full_markdown": result["full_markdown"],
            "filtered_markdown": result["filtered_markdown"],
            "page_markdowns": result["page_markdowns"],
            "page_processed_images": result["page_processed_images"],
            "json": result["json"],
            "total_pages": result["total_pages"]
        }

        logger.info(f"[{job_id}] 처리 완료 (Chandra OCR 2)")
        return JSONResponse(response)

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"[{job_id}] 내부 오류: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail="내부 서버 오류")
    finally:
        if background_tasks:
            background_tasks.add_task(cleanup_path, str(upload_path))


@app.get("/")
async def root():
    return {
        "message": "SiLIX Document Intelligence API 운영 중",
        "version": "2.0.0",
        "engine": "Chandra OCR 2",
        "endpoints": {
            "POST /process-file/": "문서 복호화 및 레이아웃 분석 (이미지 포함)",
            "Query Parameters": {
                "highqual": "bool, 고해상도 처리 여부 (기본: False)",
                "use_large_model": "bool, (호환용, Chandra에서는 무시됨)"
            }
        },
        "response_fields": {
            "full_markdown": "Base64 인코딩된 이미지 포함 마크다운",
            "filtered_markdown": "이미지 자리에 '[이미지]'만 표시된 마크다운",
            "page_markdowns": "각 페이지의 순수 텍스트 기반 마크다운 리스트",
            "page_processed_images": "각 페이지의 레이아웃 박스가 오버레이된 이미지 (Base64, PNG)"
        }
    }