import concurrent.futures
import json
from pathlib import Path

import httpx
from loguru import logger

import config

_client = httpx.Client(
    timeout=httpx.Timeout(
        connect=3.0,
        read=30.0,   # 파일 전송이므로 기존 30초 유지
        write=30.0,
        pool=3.0
    ),
    limits=httpx.Limits(
        max_keepalive_connections=5,
        keepalive_expiry=50  # A서버 --timeout-keep-alive 60 보다 약간 낮게
    )
)


def ocr_single_pdf(pdf_path: str) -> dict:
    with open(pdf_path, "rb") as f:
        resp = httpx.post(
            config.PADDLE_OCR_URL,
            files={"file": (Path(pdf_path).name, f, "application/pdf")},
            data={"group_id": "infer"},
            timeout=120,  # PDF는 이미지보다 오래 걸리므로 timeout 증가
        )
    resp.raise_for_status()
    ocr_result = resp.json()

    words, boxes = [], []
    for page in ocr_result:  # PDF는 페이지가 여러 장일 수 있으므로 전체 순회
        for text, box in zip(page.get("words", []), page.get("boxes", [])):
            text = text.strip()
            if not text or not box:
                continue
            words.append(text)
            boxes.append(box)

    logger.debug(f"PaddleOCR PDF 완료: {pdf_path} | 단어수={len(words)}")
    return {"image_path": str(pdf_path), "words": words, "boxes": boxes, "box_type": "pixel"}


def ocr_single(image_path: str, client=None) -> dict:
    with open(image_path, "rb") as f:
        data = {"group_id": "infer"}
        if config.SERVICE == "ucar":
            data["gubun"] = "ocr_with_boxes"

        resp = _client.post(
            config.PADDLE_OCR_URL,
            files={"file": (Path(image_path).name, f, "image/jpeg")},
            data=data,
        )

    resp.raise_for_status()
    ocr_result = resp.json()

    words, boxes = [], []

    # 페이지가 1장이므로 첫 번째 항목만 사용
    page = ocr_result[0] if ocr_result else {}

    for text, box in zip(page.get("words", []), page.get("boxes", [])):
        text = text.strip()
        if not text or not box:
            continue
        words.append(text)
        boxes.append(box)

    result = {"image_path": str(image_path), "words": words, "boxes": boxes, "box_type": "pixel"}

    # OCR 결과 저장
    # json_path = Path(image_path).with_name(Path(image_path).stem + "_paddle.json")
    # with open(json_path, "w", encoding="utf-8") as f:
    #     json.dump(result, f, ensure_ascii=False, indent=2)
    # logger.debug(f"OCR 저장: {json_path}")

    logger.debug(f"PaddleOCR 완료: {image_path} | 단어수={len(words)}")
    return result


def ocr_batch(image_paths: list, client=None, max_workers: int = config.OCR_MAX_WORKERS,
              on_done=None) -> list:   # ← on_done 추가
    results = [None] * len(image_paths)

    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_idx = {
            executor.submit(ocr_single, str(p)): i
            # executor.submit(ocr_single_pdf, str(p)): i
            for i, p in enumerate(image_paths)
        }
        for future in concurrent.futures.as_completed(future_to_idx):
            idx = future_to_idx[future]
            try:
                results[idx] = future.result()
            except Exception as e:
                logger.error(f"PaddleOCR 실패: {image_paths[idx]} | {e}")
                results[idx] = {"image_path": str(image_paths[idx]),
                                "words": [], "boxes": [], "error": str(e)}
            if on_done:          # ← 콜백 호출
                on_done(idx, results[idx])
    return results