# ============================================================================
# videocr(PaddleOCR) 기반 "영상 박힌 자막 → SRT/세그먼트" FastAPI 엔드포인트 예시
# h-python 서비스(기존 /transcribe·/ocr 있는 FastAPI 앱)에 그대로 붙일 수 있게 작성.
#
# 응답 형식을 기존 /transcribe 와 동일하게 맞춤:
#   { "language": "korean", "segments": [ {"start":0.0,"end":2.7,"text":"..."} ], "srt": "..." }
# → h-lab 쪽은 기존 전사 저장 로직을 거의 그대로 재사용 가능.
#
# ── 설치 ────────────────────────────────────────────────────────────────────
#   pip install videocr-PaddleOCR paddleocr
#   # CPU: pip install paddlepaddle
#   # GPU: pip install paddlepaddle-gpu   (CUDA 환경일 때만, 훨씬 빠름)
#
# ── 언어 코드(PaddleOCR) ─────────────────────────────────────────────────────
#   korean / japan / en / ch / chinese_cht ...
#   (h-lab는 ko/ja/en/zh 로 보낼 수 있으니 아래 LANG_MAP 에서 변환)
# ============================================================================

import os
import re
import tempfile

from fastapi import APIRouter, UploadFile, File, Form, HTTPException

# videocr-PaddleOCR
from videocr import get_subtitles

router = APIRouter()  # 기존 app 에 app.include_router(router) 하거나, 아래 데코를 @app.post 로 바꿔 사용

# h-lab(ko/ja/en/zh) → PaddleOCR 언어 코드
LANG_MAP = {
    "ko": "korean", "kor": "korean", "korean": "korean",
    "ja": "japan", "jpn": "japan", "japan": "japan",
    "en": "en", "eng": "en",
    "zh": "ch", "chi_sim": "ch", "ch": "ch",
}


def _to_seconds(ts: str) -> float:
    """'HH:MM:SS,mmm' → 초(float)."""
    ts = ts.strip().replace(".", ",")
    h, m, rest = ts.split(":")
    s, ms = (rest.split(",") + ["0"])[:2]
    return int(h) * 3600 + int(m) * 60 + int(s) + int(ms.ljust(3, "0")[:3]) / 1000.0


def _parse_srt(srt_text: str):
    """SRT 문자열 → [{start, end, text}] 세그먼트 리스트."""
    segments = []
    # 블록: (선택)인덱스줄 / 시간줄 / 텍스트(여러 줄)
    blocks = re.split(r"\n\s*\n", srt_text.strip())
    time_re = re.compile(r"(\d{1,2}:\d{2}:\d{2}[,.]\d{1,3})\s*-->\s*(\d{1,2}:\d{2}:\d{2}[,.]\d{1,3})")
    for block in blocks:
        lines = [ln for ln in block.splitlines() if ln.strip() != ""]
        if not lines:
            continue
        # 시간줄 찾기(인덱스줄이 있을 수도/없을 수도)
        time_idx = next((i for i, ln in enumerate(lines) if time_re.search(ln)), None)
        if time_idx is None:
            continue
        m = time_re.search(lines[time_idx])
        text = " ".join(ln.strip() for ln in lines[time_idx + 1:]).strip()
        if not text:
            continue
        segments.append({
            "start": round(_to_seconds(m.group(1)), 3),
            "end": round(_to_seconds(m.group(2)), 3),
            "text": text,
        })
    return segments


@router.post("/ocr_video")
async def ocr_video(
    file: UploadFile = File(..., description="영상 파일(mp4 등)"),
    lang: str = Form("ko", description="ko|ja|en|zh (PaddleOCR 코드로 자동 변환)"),
    # 자막 영역만 OCR 하면 정확도·속도 둘 다 좋아짐. 비우면 하단 위주(use_fullframe=False).
    use_fullframe: bool = Form(False, description="True면 전체 프레임 OCR(느림), False면 하단 자막영역 위주"),
    crop_x: int = Form(None), crop_y: int = Form(None),
    crop_width: int = Form(None), crop_height: int = Form(None),
    # 성능/정확도 튜닝
    frames_to_skip: int = Form(1, description="프레임 N개 건너뛰며 OCR(클수록 빠름·덜 정밀)"),
    conf_threshold: int = Form(75, description="OCR 신뢰도 임계(0~100)"),
    sim_threshold: int = Form(80, description="연속 자막 병합 유사도(0~100)"),
    time_start: str = Form("0:00"), time_end: str = Form(""),
    use_gpu: bool = Form(False),
):
    """영상에 박힌(하드코딩) 자막을 OCR로 추출해 SRT/세그먼트로 반환한다."""
    paddle_lang = LANG_MAP.get((lang or "ko").lower(), "korean")

    # 업로드 영상을 임시파일로 저장(videocr 는 파일 경로 필요)
    suffix = os.path.splitext(file.filename or "")[1] or ".mp4"
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
    try:
        tmp.write(await file.read())
        tmp.flush()
        tmp.close()

        # videocr: 영상 → SRT 문자열 (프레임 샘플링 + OCR + 연속자막 병합 내장)
        srt_text = get_subtitles(
            tmp.name,
            lang=paddle_lang,
            use_gpu=use_gpu,
            time_start=time_start,
            time_end=time_end,
            conf_threshold=conf_threshold,
            sim_threshold=sim_threshold,
            use_fullframe=use_fullframe,
            frames_to_skip=frames_to_skip,
            crop_x=crop_x, crop_y=crop_y,
            crop_width=crop_width, crop_height=crop_height,
        )

        segments = _parse_srt(srt_text)
        return {
            "language": paddle_lang,
            "segments": segments,
            "srt": srt_text,
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"OCR 영상 처리 실패: {e}")
    finally:
        try:
            os.unlink(tmp.name)
        except OSError:
            pass


# ── 기존 app 에 붙이는 방법 ───────────────────────────────────────────────────
# from fastapi import FastAPI
# app = FastAPI()
# app.include_router(router)            # → POST /ocr_video
#
# 또는 router 안 쓰고 직접:
#   @app.post("/ocr_video")
#   async def ocr_video(...): ...
#
# ── 빠른 테스트 ──────────────────────────────────────────────────────────────
#   curl -X POST http://localhost:8000/ocr_video \
#        -F "file=@shorts.mp4" -F "lang=ja" -F "frames_to_skip=2"
#   → {"language":"japan","segments":[{"start":0.0,"end":2.7,"text":"..."}],"srt":"..."}