"""자막 교정 — Google Gemini(무료 티어) 사용. 텍스트만 전송(영상 X).

키 없으면 교정 건너뜀(들리는대로 그대로). 줄 수/순서 보존, 사투리·구어체 유지.
"""
from __future__ import annotations

import base64
import json
import os
import urllib.error
import urllib.request
from typing import List, Optional, Tuple


class GeminiQuotaError(Exception):
    """Gemini 무료 한도 초과(429) — Whisper 폴백 신호."""

GEMINI_MODEL = "gemini-2.5-flash"  # 무료 티어 지원(2.0-flash는 limit 0)
_ENDPOINT = "https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={key}"

_PROMPT = """너는 한국어 영상 자막 교정기다. 아래 자막 줄들을 교정해라.

규칙:
- 명백한 오타·오인식, 고유명사(인명/지명/유행어), 띄어쓰기만 고친다.
- 사투리·반말·구어체 말투는 그대로 둔다(표준어로 바꾸지 마라).
- 외국 문자나 깨진 글자는 들리는 한국어로 자연스럽게 바꾼다.
- 줄 수와 순서를 절대 바꾸지 마라. 각 줄을 1:1로 교정. 합치거나 나누지 마라.
- 도저히 못 고치겠으면 원문 그대로 둔다.
영상 제목(문맥): {title}

자막(줄 순서대로):
{lines}

교정된 줄들을 JSON 문자열 배열로만 출력. 입력과 같은 개수."""


def _gemini_key() -> Optional[str]:
    key = os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
    if key:
        return key.strip()
    # 프로젝트 루트의 .gemini_key 파일 폴백
    root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    f = os.path.join(root, ".gemini_key")
    if os.path.isfile(f):
        return open(f, encoding="utf-8").read().strip() or None
    return None


def correct_captions(texts: List[str], *, title: str = "",
                     model: str = GEMINI_MODEL, key: Optional[str] = None) -> List[str]:
    """자막 리스트를 Gemini로 교정. 실패/키없음/개수불일치 시 원문 반환(안전)."""
    key = key or _gemini_key()
    if not key or not texts:
        return texts

    numbered = "\n".join(f"{i+1}. {t}" for i, t in enumerate(texts))
    body = {
        "contents": [{"parts": [{"text": _PROMPT.format(title=title or "(없음)", lines=numbered)}]}],
        "generationConfig": {
            "temperature": 0.2,
            "responseMimeType": "application/json",
            "responseSchema": {"type": "array", "items": {"type": "string"}},
        },
    }
    url = _ENDPOINT.format(model=model, key=key)
    req = urllib.request.Request(
        url, data=json.dumps(body).encode("utf-8"),
        headers={"Content-Type": "application/json"},
    )
    try:
        with urllib.request.urlopen(req, timeout=60) as resp:
            data = json.loads(resp.read().decode("utf-8"))
        out = data["candidates"][0]["content"]["parts"][0]["text"]
        fixed = json.loads(out)
    except (urllib.error.URLError, KeyError, IndexError, json.JSONDecodeError, TimeoutError):
        return texts  # 안전: 실패 시 원문

    if not isinstance(fixed, list) or len(fixed) != len(texts):
        return texts  # 개수 안 맞으면 매핑 깨지니 원문
    return [str(f).strip() or texts[i] for i, f in enumerate(fixed)]


def has_gemini_key() -> bool:
    return bool(_gemini_key())


_ASR_PROMPT = (
    "이 한국어 영상 오디오를 받아쓰기 해줘. 말소리가 있는 모든 구간을 빠짐없이.\n"
    "자막용으로 짧게 나눠라:\n"
    "- 한 자막은 한국어 12자 내외(최대 14자).\n"
    "- 의미가 자연스럽게 끊기는 지점(어절·구·절 경계)에서 나눠라. 단어 중간이나 "
    "어색한 곳(조사 앞 등)에서 끊지 마라. 짧은 한 호흡이 한 자막.\n"
    "각 자막을 {\"start\":초, \"end\":초, \"text\":\"...\"} 로(시간은 오디오 시작 기준 초, 소수1자리). "
    "사투리·반말·구어체는 그대로. 웃음/박수 등 비언어는 제외. JSON 배열로만."
)


def transcribe_gemini(audio_path: str, *, model: str = GEMINI_MODEL,
                      key: Optional[str] = None) -> List[Tuple[float, float, str]]:
    """Gemini로 오디오 받아쓰기 → [(start, end, text)] (오디오 시작 기준 초).

    429(무료 한도 초과) 시 GeminiQuotaError. 그 외 실패는 RuntimeError.
    """
    key = key or _gemini_key()
    if not key:
        raise RuntimeError("Gemini 키 없음")
    audio_b64 = base64.b64encode(open(audio_path, "rb").read()).decode()
    body = {
        "contents": [{"parts": [
            {"inline_data": {"mime_type": "audio/mp3", "data": audio_b64}},
            {"text": _ASR_PROMPT},
        ]}],
        "generationConfig": {
            "temperature": 0.1, "responseMimeType": "application/json",
            "responseSchema": {"type": "array", "items": {"type": "object", "properties": {
                "start": {"type": "number"}, "end": {"type": "number"}, "text": {"type": "string"}},
                "required": ["start", "end", "text"]}},
        },
    }
    url = _ENDPOINT.format(model=model, key=key)
    req = urllib.request.Request(url, data=json.dumps(body).encode("utf-8"),
                                 headers={"Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=180) as resp:
            data = json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        if e.code == 429:
            raise GeminiQuotaError()
        raise RuntimeError(f"Gemini HTTP {e.code}")
    out = json.loads(data["candidates"][0]["content"]["parts"][0]["text"])
    segs = []
    for s in out:
        try:
            st, en, tx = float(s["start"]), float(s["end"]), str(s["text"]).strip()
        except (KeyError, ValueError, TypeError):
            continue
        if tx and en > st:
            segs.append((st, en, tx))
    return segs