capcut-agent/capcut_agent/plan.py

"""자동 탭 — Gemini로 오팔 파이프라인 대체.

Step 1: 전체 영상(저 fps) → 하이라이트 후보 5개.
Step 3: 구간별(기본 fps, videoMetadata로 클립) → 붙여넣기 탭 스키마 JSON + 타이틀 후보 5선.
오팔 Step 2(노드 간 구간 값 전달)는 코드에선 함수 인자이므로 존재하지 않는다.
"""
from __future__ import annotations

import json
import re
import urllib.error
import urllib.request
from typing import Dict, List, Optional

from .correct import GeminiQuotaError, _gemini_key
from .paste import parse_paste, parse_time
from . import prompts

DEFAULT_MODEL = "gemini-3.5-flash"
_ENDPOINT = "https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={key}"

_FENCE_RE = re.compile(r"```json\s*(.*?)```", re.S)
# 블록 ② "1. 상단: … / 메인: …  — [유형]" — 번호·대괄호는 있어도 없어도, 대시 3종 허용
_TITLE_RE = re.compile(
    r"^\s*(?:\d+\.\s*)?상단\s*[:：]\s*(.+?)\s*/\s*메인\s*[:：]\s*(.+?)(?:\s*[—–-]\s*(.+?))?\s*$",
    re.M)


def parse_titles(text: str) -> List[Dict]:
    """텍스트에서 타이틀 후보(상단/메인/유형) 추출. 오팔 붙여넣기·Step3 응답 공용."""
    out: List[Dict] = []
    for m in _TITLE_RE.finditer(text or ""):
        out.append({"top": m.group(1).strip(), "main": m.group(2).strip(),
                    "kind": (m.group(3) or "").strip().strip("[]")})
    return out


def _fmt_offset(sec: float) -> str:
    return f"{max(0, int(sec))}s"


def _call(video_url: str, prompt_text: str, *, model: str, key: str,
          fps: Optional[float] = None, start: Optional[float] = None,
          end: Optional[float] = None, timeout: float = 600.0) -> str:
    """Gemini generateContent 1회 → 응답 텍스트. 429는 GeminiQuotaError."""
    part: Dict = {"fileData": {"fileUri": video_url}}
    meta: Dict = {}
    if fps is not None and abs(fps - 1.0) > 1e-9:
        meta["fps"] = fps
    if start is not None and end is not None:
        meta["startOffset"] = _fmt_offset(start)
        meta["endOffset"] = _fmt_offset(end)
    if meta:
        part["videoMetadata"] = meta
    body = {
        "contents": [{"parts": [part, {"text": prompt_text}]}],
        "generationConfig": {"temperature": 0.7},
    }
    req = urllib.request.Request(
        _ENDPOINT.format(model=model, key=key),
        data=json.dumps(body).encode("utf-8"),
        headers={"Content-Type": "application/json"},
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            data = json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        if e.code == 429:
            raise GeminiQuotaError()
        try:
            detail = e.read().decode("utf-8", "replace")[:400]
        except Exception:  # noqa: BLE001
            detail = ""
        raise RuntimeError(f"Gemini HTTP {e.code}: {detail}")
    try:
        return data["candidates"][0]["content"]["parts"][0]["text"]
    except (KeyError, IndexError, TypeError):
        raise RuntimeError(
            "Gemini 응답 형식 오류: "
            + json.dumps(data, ensure_ascii=False)[:300])


def _extract_json_str(text: str) -> str:
    """첫 ```json 펜스 안쪽. 펜스 없으면 가장 바깥 {…} 범위, 그것도 없으면 전문."""
    m = _FENCE_RE.search(text or "")
    if m:
        return m.group(1).strip()
    t = (text or "").strip()
    i, j = t.find("{"), t.rfind("}")
    if i != -1 and j > i:
        return t[i:j + 1]
    return t


def parse_candidates(text: str, *, src: str = "Step 1 응답") -> List[Dict]:
    """`{"candidates":[{id,start_time,end_time,reason}, …]}` → [{id,start,end,reason}] (초).

    Gemini Step1 응답과 사용자가 직접 붙여넣는 구간 JSON이 **같은 형식**이라 둘이 공유한다.
    코드펜스(```)·앞뒤 잡텍스트는 `_extract_json_str`이 걷어낸다.
    """
    try:
        data = json.loads(_extract_json_str(text), strict=False)
    except json.JSONDecodeError:
        raise RuntimeError(f"{src} JSON 파싱 실패:\n{text[:400]}")
    cands = data.get("candidates") if isinstance(data, dict) else None
    if not isinstance(cands, list) or not cands:
        raise RuntimeError(f"{src}에 candidates 배열이 없습니다:\n{text[:400]}")
    out: List[Dict] = []
    for i, c in enumerate(cands, 1):
        if not isinstance(c, dict):
            continue
        try:
            s = parse_time(c.get("start_time"))
            e = parse_time(c.get("end_time"))
        except (ValueError, TypeError):
            continue
        if e > s:
            out.append({"id": int(c.get("id") or i), "start": s, "end": e,
                        "reason": str(c.get("reason") or "").strip()})
    if not out:
        raise RuntimeError(f"{src}: 유효한 구간이 하나도 없습니다.")
    return out


def select_highlights(url: str, *, key: Optional[str] = None,
                      model: Optional[str] = None) -> List[Dict]:
    """Step 1: 전체 영상 → 하이라이트 후보 [{id, start, end, reason}] (초)."""
    key = key or _gemini_key()
    if not key:
        raise RuntimeError("Gemini 키 없음 (.gemini_key)")
    cfg = prompts.load_config()
    model = model or cfg.get("model_step1") or cfg.get("model") or DEFAULT_MODEL
    fps = float(cfg.get("fps_step1") or 0.2)
    text = _call(url, prompts.load_step1(), model=model, key=key, fps=fps)
    return parse_candidates(text)


def _shift_cuts_if_relative(payload: Dict, start: float, end: float) -> str:
    """클립 기준(0부터) 타임코드로 보이면 원본 기준으로 보정. 반환: 로그용 메모.

    구간을 잘라 보낸 클립에 대해 모델이 절대/클립 어느 기준으로 답할지 보장이
    없다(문서 미명시). 절대 기준으로 보이면 그대로(우선), 클립 기준으로 보이면
    start 를 더하고, 둘 다 아니면 손대지 않는다(이후 검증·다운로드에서 드러남).
    """
    cuts = payload["cuts"]
    dur, pad = end - start, 10.0
    if all(start - pad <= s and e <= end + pad for s, e, _, _ in cuts):
        return ""                                    # 절대 기준 — 그대로
    if start > pad and all(0 <= s and e <= dur + pad for s, e, _, _ in cuts):
        payload["cuts"] = [(s + start, e + start, b, f) for s, e, b, f in cuts]
        return f"컷 타임코드가 클립 기준 → +{start:.0f}s 보정"
    return ""


def edit_plan(url: str, start: float, end: float, *, key: Optional[str] = None,
              model: Optional[str] = None) -> Dict:
    """Step 3: [start, end] 클립 → {paste, titles, time_note}.

    paste 는 기존 paste.parse_paste 검증을 그대로 통과한 결과이며 url 은 입력값으로
    덮어쓴다(LLM이 영상 ID를 지어내는 사고 차단).
    """
    key = key or _gemini_key()
    if not key:
        raise RuntimeError("Gemini 키 없음 (.gemini_key)")
    cfg = prompts.load_config()
    model = model or cfg.get("model") or DEFAULT_MODEL
    fps = float(cfg.get("fps_step3") or 1.0)
    text = _call(url, prompts.load_step3(), model=model, key=key,
                 fps=fps, start=start, end=end)
    # LLM이 쓴 url 은 신뢰하지 않는다 — 검증(parse_paste) 전에 입력 URL 로 강제 교체.
    # (Gemini가 가끔 url 을 지어내거나 깨뜨려 ValueError 로 하이라이트가 통째로 죽는다)
    try:
        raw = json.loads(_extract_json_str(text), strict=False)
    except json.JSONDecodeError as e:
        raise ValueError(f"JSON 형식 오류: {e.msg} (줄 {e.lineno})")
    if isinstance(raw, dict):
        raw["url"] = url
    payload = parse_paste(raw)                       # ValueError 는 호출부에서 표시
    payload["url"] = url
    note = _shift_cuts_if_relative(payload, start, end)
    return {"paste": payload, "titles": parse_titles(text), "time_note": note}