utils.py

from __future__ import annotations

import os
import random
import re
import shutil
from pathlib import Path

import pysubs2
from faster_whisper import WhisperModel

# Résolution de référence du fichier ASS (identique à la vidéo de sortie 9:16)
_ASS_PLAY_RES_X = 1080
_ASS_PLAY_RES_Y = 1920

# Police embarquée pour les sous-titres (chargée par ffmpeg via fontsdir)
ASSETS_DIR = Path(__file__).resolve().parent / "assets"
FONTS_DIR = ASSETS_DIR / "fonts"
CAPTION_FONT = "Montserrat ExtraBold" if any(FONTS_DIR.glob("*.ttf")) else "Arial Black"

# Musiques de fond (mixées à bas volume avec ducking sous la voix)
MUSIC_DIR = ASSETS_DIR / "music"
_MUSIC_EXTS = {".mp3", ".m4a", ".aac", ".wav", ".flac", ".ogg"}

# Watermark incrustée en bas à droite de la vidéo finale
WATERMARK_DIR = ASSETS_DIR / "watermark"
_WATERMARK_EXTS = {".png", ".jpg", ".jpeg", ".webp"}


def pick_music_track() -> str | None:
    """
    Choisit aléatoirement une musique de fond dans assets/music/.
    Retourne None si le dossier est absent ou vide → pas de musique.
    """
    if not MUSIC_DIR.is_dir():
        return None
    tracks = [
        p for p in MUSIC_DIR.iterdir()
        if p.suffix.lower() in _MUSIC_EXTS and not p.name.startswith(".")
    ]
    return str(random.choice(tracks)) if tracks else None


def pick_watermark(dark_video: bool | None = None) -> str | None:
    """
    Choisit la watermark dans assets/watermark/ selon la luminosité de la
    vidéo : `white-watermark.png` (logo blanc) si la vidéo est sombre,
    `black-watermark.png` (logo noir) si elle est claire — pour un contraste
    maximal. Si la variante attendue est absente (ou `dark_video` inconnu),
    retombe sur la première image du dossier. Retourne None si le dossier est
    absent ou vide → pas de watermark.

    :param dark_video: True si la vidéo est sombre, None si non mesurée.
    """
    if not WATERMARK_DIR.is_dir():
        return None
    images = sorted(
        p for p in WATERMARK_DIR.iterdir()
        if p.suffix.lower() in _WATERMARK_EXTS and not p.name.startswith(".")
    )
    if dark_video is not None:
        wanted = "white-watermark" if dark_video else "black-watermark"
        for p in images:
            if p.stem.lower() == wanted:
                return str(p)
    return str(images[0]) if images else None

# Nombre max de mots affichés simultanément (style CapCut)
_WORDS_PER_CHUNK = 3

_whisper_model: WhisperModel | None = None


def _get_whisper_model() -> WhisperModel:
    """Charge le modèle Whisper une seule fois (singleton paresseux)."""
    global _whisper_model
    model = _whisper_model
    if model is None:
        model = WhisperModel("base", device="cpu", compute_type="int8")
        _whisper_model = model
    return model


def get_downloads_dir() -> Path:
    """Return the user's Downloads folder, falling back to creating it."""
    home = Path(os.path.expanduser("~"))
    for name in ("Downloads", "T\xe9l\xe9chargements"):
        candidate = home / name
        if candidate.is_dir():
            return candidate
    downloads = home / "Downloads"
    downloads.mkdir(parents=True, exist_ok=True)
    return downloads


def safe_rmtree(path: str | Path) -> None:
    """Remove a directory tree without raising if it fails."""
    try:
        shutil.rmtree(path)
    except OSError:
        pass


def generate_styled_subtitles(audio_path: str) -> tuple[str, list[dict]]:
    """
    Transcrit un fichier audio avec faster-whisper et génère des sous-titres
    style "CapCut" : 1 à 3 mots affichés à la fois, centrés, en majuscules,
    mot actif en jaune avec un effet pop, synchronisés sur les timestamps réels.

    Fonction synchrone et bloquante : l'appelant la lance dans un thread
    (asyncio.to_thread) pour la paralléliser avec le traitement vidéo.

    :param audio_path: Chemin vers le fichier audio.
    :return: (chemin du fichier ASS, liste des phrases {start, end, text})
             — les phrases servent à aligner les coupes de clips.
    """
    model = _get_whisper_model()
    segments, _info = model.transcribe(audio_path, word_timestamps=True)

    subs = pysubs2.SSAFile()
    subs.info["PlayResX"] = str(_ASS_PLAY_RES_X)
    subs.info["PlayResY"] = str(_ASS_PLAY_RES_Y)

    style = subs.styles["Default"]
    style.fontname = CAPTION_FONT
    style.fontsize = 110
    style.bold = True
    style.primarycolor = pysubs2.Color(255, 255, 255, 0)
    style.outlinecolor = pysubs2.Color(0, 0, 0, 0)
    style.outline = 5
    style.shadow = 2
    style.alignment = pysubs2.Alignment.BOTTOM_CENTER
    style.marginl = 60
    style.marginr = 60
    style.marginv = 700  # ~centre-bas, au-dessus de la zone UI TikTok

    sentences: list[dict] = []
    for seg in segments:
        words = seg.words or []
        if not words:
            continue
        sentences.append(
            {
                "start": float(words[0].start),
                "end": float(words[-1].end),
                "text": (seg.text or "").strip(),
            }
        )

        for i in range(0, len(words), _WORDS_PER_CHUNK):
            chunk = words[i:i + _WORDS_PER_CHUNK]
            chunk_end = float(chunk[-1].end)
            # Un event par mot : le chunk entier est affiché, le mot courant
            # passe en jaune, les autres restent blancs. Pas d'animation de
            # taille : un mot qui grossit décale toute la ligne centrée et
            # fait "trembler" le texte.
            for j, word in enumerate(chunk):
                ev_start = float(word.start)
                ev_end = float(chunk[j + 1].start) if j + 1 < len(chunk) else chunk_end
                if ev_end <= ev_start:
                    ev_end = ev_start + 0.05
                parts = []
                for k, other in enumerate(chunk):
                    text = other.word.strip().upper()
                    if k == j:
                        parts.append(r"{\c&H00FFFF&}" + text + r"{\r}")
                    else:
                        parts.append(text)
                subs.events.append(
                    pysubs2.SSAEvent(
                        start=int(ev_start * 1000),
                        end=int(ev_end * 1000),
                        text=" ".join(parts),
                    )
                )

    ass_path = os.path.join(Path(audio_path).parent, "subtitles.ass")
    subs.save(ass_path)
    return ass_path, sentences


def add_hook_events(ass_path: str, hooks: list[tuple[float, float, str]]) -> None:
    """
    Incruste un texte "hook" (accroche) en haut de l'écran pendant les
    premières secondes de chaque clip. Modifie le fichier ASS en place.
    Pas de fondu d'entrée : le hook doit être visible dès la toute première
    frame (c'est elle qui sert de couverture TikTok).

    :param hooks: liste de (start_s, end_s, texte).
    """
    subs = pysubs2.load(ass_path)

    hook_style = subs.styles["Default"].copy()
    hook_style.fontsize = 84
    hook_style.primarycolor = pysubs2.Color(255, 255, 0, 0)  # jaune
    hook_style.alignment = pysubs2.Alignment.TOP_CENTER
    hook_style.marginv = 420
    subs.styles["Hook"] = hook_style

    for start, end, text in hooks:
        text = (text or "").strip().upper()
        if not text or end <= start:
            continue
        if len(text) > 70:
            text = text[:67] + "…"
        subs.events.append(
            pysubs2.SSAEvent(
                start=int(start * 1000),
                end=int(end * 1000),
                text=r"{\fad(0,150)}" + text,
                style="Hook",
                layer=1,
            )
        )
    subs.save(ass_path)


# ─────────────────────────────────────────────────────────────────────────────
# Jump cuts : suppression des silences avec recalage de timeline
# ─────────────────────────────────────────────────────────────────────────────
class TimeMap:
    """
    Convertit un timestamp de la timeline originale vers la timeline obtenue
    après suppression des silences (segments conservés concaténés bout à bout).
    """

    def __init__(self, keep_segments: list[tuple[float, float]]):
        self.segments = sorted(keep_segments)
        self._new_starts: list[float] = []
        acc = 0.0
        for a, b in self.segments:
            self._new_starts.append(acc)
            acc += b - a
        self.new_total = acc

    def map(self, t: float) -> float:
        """Timestamp original → timestamp après coupes (clampé aux bornes)."""
        for (a, b), new_start in zip(self.segments, self._new_starts):
            if t < a:
                return new_start  # dans un silence coupé → collé au segment suivant
            if t <= b:
                return new_start + (t - a)
        return self.new_total

    def unmap(self, t: float) -> float:
        """Timestamp après coupes → timestamp original (inverse de map)."""
        for (a, b), new_start in zip(self.segments, self._new_starts):
            if t <= new_start + (b - a):
                return a + max(0.0, t - new_start)
        return self.segments[-1][1] if self.segments else t


def compute_keep_segments(
    sentences: list[dict],
    total: float,
    min_gap: float = 1.0,
    pad: float = 0.25,
    max_segments: int = 120,
    fps: float = 0.0,
) -> list[tuple[float, float]] | None:
    """
    Calcule les segments de parole à CONSERVER pour les jump cuts : les
    silences entre phrases ≥ min_gap sont supprimés, avec `pad` secondes de
    respiration autour de chaque zone de parole.

    Garde-fous : retourne None (= ne rien couper) si le transcript est vide,
    si la parole couvre < 40 % de la vidéo (musique, transcription ratée) ou
    s'il n'y a presque rien à couper (> 98 % de couverture).

    :param fps: si fournie, les bornes sont alignées sur la grille des frames
                (entre deux frames) pour garder audio et vidéo exactement synchrones.
    """
    if not sentences or total <= 0:
        return None

    speech = sorted((max(0.0, s["start"]), min(total, s["end"])) for s in sentences)
    merged: list[list[float]] = []
    for a, b in speech:
        a, b = max(0.0, a - pad), min(total, b + pad)
        if merged and a - merged[-1][1] < min_gap:
            merged[-1][1] = max(merged[-1][1], b)
        else:
            merged.append([a, b])

    coverage = sum(b - a for a, b in merged) / total
    if coverage < 0.40 or coverage > 0.98:
        return None

    # Trop de coupes → ne garder que les plus gros silences (l'expression
    # ffmpeg select doit rester de taille raisonnable)
    if len(merged) > max_segments:
        biggest_gaps = set(
            sorted(
                range(1, len(merged)),
                key=lambda i: merged[i][0] - merged[i - 1][1],
                reverse=True,
            )[: max_segments - 1]
        )
        rebuilt = [merged[0][:]]
        for i in range(1, len(merged)):
            if i in biggest_gaps:
                rebuilt.append(merged[i][:])
            else:
                rebuilt[-1][1] = merged[i][1]
        merged = rebuilt

    # Aligner les bornes entre deux frames → durées vidéo/audio identiques
    if fps and fps > 0:
        snapped = []
        for a, b in merged:
            a = max(0.0, (round(a * fps) - 0.5) / fps)
            b = min(total, (round(b * fps) - 0.5) / fps)
            if b - a > 0.2:
                snapped.append([a, b])
        merged = snapped
        if not merged:
            return None

    return [(a, b) for a, b in merged]


def restrict_to_windows(
    windows: list[dict],
    tmap: TimeMap,
    fps: float = 0.0,
) -> tuple[list[tuple[float, float]], list[dict]]:
    """
    Restreint le rendu au seul contenu des clips : croise chaque fenêtre
    (timeline post-silences) avec les segments de parole conservés, et
    reconstruit la timeline finale où les clips sont mis bout à bout.
    Tout ce qui n'appartient à aucun clip ne sera ni encodé ni décodé en sortie.

    :param windows: fenêtres de clips sur la timeline de `tmap` (post-coupes).
    :param tmap: TimeMap des coupes de silences (timeline originale → post-coupes).
    :param fps: alignement des bornes sur la grille des frames (synchro A/V).
    :return: (segments originaux à encoder, fenêtres recalées sur la timeline finale).
    """
    keep: list[tuple[float, float]] = []
    new_windows: list[dict] = []
    acc = 0.0
    for w in windows:
        orig_start, orig_end = tmap.unmap(w["start"]), tmap.unmap(w["end"])
        segs: list[tuple[float, float]] = []
        for a, b in tmap.segments:
            lo, hi = max(a, orig_start), min(b, orig_end)
            if hi - lo <= 1e-3:
                continue
            if fps and fps > 0:
                lo = max(0.0, (round(lo * fps) - 0.5) / fps)
                hi = (round(hi * fps) - 0.5) / fps
            if hi - lo > 0.05:
                segs.append((lo, hi))
        duration = sum(b - a for a, b in segs)
        if duration <= 0.05:
            continue
        keep.extend(segs)
        new_windows.append({**w, "start": acc, "end": acc + duration})
        acc += duration

    # Fusionner les segments qui se touchent (fenêtres adjacentes)
    merged: list[list[float]] = []
    for a, b in keep:
        if merged and a - merged[-1][1] < 0.02:
            merged[-1][1] = max(merged[-1][1], b)
        else:
            merged.append([a, b])
    return [(a, b) for a, b in merged], new_windows


def remap_sentences(sentences: list[dict], tmap: TimeMap) -> list[dict]:
    """Recale les phrases sur la timeline post-coupes (jump cuts)."""
    out = []
    for s in sentences:
        ns, ne = tmap.map(s["start"]), tmap.map(s["end"])
        if ne - ns > 0.05:
            out.append({**s, "start": ns, "end": ne})
    return out


def remap_ass(ass_path: str, tmap: TimeMap) -> None:
    """Recale tous les événements du fichier ASS sur la timeline post-coupes."""
    subs = pysubs2.load(ass_path)
    kept = []
    for ev in subs.events:
        ns = tmap.map(ev.start / 1000.0)
        ne = tmap.map(ev.end / 1000.0)
        if ne - ns > 0.03:
            ev.start, ev.end = int(ns * 1000), int(ne * 1000)
            kept.append(ev)
    subs.events = kept
    subs.save(ass_path)


# ─────────────────────────────────────────────────────────────────────────────
# Surlignage des mots-clés dans les sous-titres
# ─────────────────────────────────────────────────────────────────────────────
_KEYWORD_TAG = r"{\c&H66FF00&}"  # vert fluo (format ASS &HBBGGRR&)


def _normalize_word(word: str) -> str:
    return re.sub(r"[^\w]", "", word, flags=re.UNICODE).casefold()


def highlight_keywords(ass_path: str, windows: list[dict]) -> None:
    """
    Surligne en vert, dans les sous-titres, les mots-clés choisis par le LLM
    pour chaque clip. Le mot actif (jaune + pop) garde la priorité.
    Modifie le fichier ASS en place.
    """
    spans = [
        (w["start"], w["end"], {_normalize_word(k) for k in w.get("keywords") or []})
        for w in windows
        if w.get("keywords")
    ]
    if not spans:
        return

    subs = pysubs2.load(ass_path)
    for ev in subs.events:
        if ev.style != "Default":
            continue
        mid = (ev.start + ev.end) / 2000.0
        keywords = next((k for a, b, k in spans if a <= mid <= b), None)
        if not keywords:
            continue
        tokens = ev.text.split(" ")
        changed = False
        for idx, tok in enumerate(tokens):
            if "{" in tok:  # mot actif déjà stylé → ne pas toucher
                continue
            if _normalize_word(tok) in keywords:
                tokens[idx] = _KEYWORD_TAG + tok + r"{\r}"
                changed = True
        if changed:
            ev.text = " ".join(tokens)
    subs.save(ass_path)


def select_best_clips(
    sentences: list[dict],
    total_duration: float,
    max_clips: int = 5,
    context: str = "",
    log=None,
) -> list[dict] | None:
    """
    Sélectionne les meilleurs passages du transcript via l'API Claude et
    génère hook/description/hashtags/mots-clés/score de viralité par clip
    (sortie JSON structurée).

    Nécessite ANTHROPIC_API_KEY dans l'environnement. Retourne None si la clé
    est absente ou en cas d'erreur → l'appelant bascule sur le découpage
    séquentiel aligné sur les phrases.

    :param context: contexte optionnel (titre/auteur de la vidéo) pour aider la sélection.
    :return: liste de {start, end, hook, description, hashtags, keywords,
             virality_score} triée par score décroissant, ou None.
    """
    if not os.environ.get("ANTHROPIC_API_KEY") or not sentences:
        return None

    def _log(msg: str) -> None:
        if log:
            log(msg)

    try:
        import anthropic
        from pydantic import BaseModel

        class ClipPlan(BaseModel):
            first_sentence: int
            last_sentence: int
            hook: str
            description: str
            hashtags: list[str]
            keywords: list[str]
            virality_score: int

        class ClipSelection(BaseModel):
            clips: list[ClipPlan]

        transcript = "\n".join(
            f"{i}|{s['start']:.0f}s|{s['text']}" for i, s in enumerate(sentences)
        )
        context_block = f"Contexte de la vidéo : {context}\n" if context else ""
        prompt = (
            "Voici le transcript d'une vidéo, une phrase par ligne au format "
            "index|début_en_secondes|texte.\n"
            f"{context_block}"
            f"Sélectionne les {max_clips} meilleurs passages pour en faire des clips "
            "TikTok viraux et autonomes (compréhensibles sans le reste de la vidéo).\n"
            "Contraintes :\n"
            "- Un clip = un intervalle de phrases consécutives, désigné par les index "
            "first_sentence et last_sentence (inclus).\n"
            "- Chaque clip dure entre 15 et 60 secondes.\n"
            "- Les clips ne se chevauchent pas.\n"
            "- hook : accroche percutante de 8 mots max, impérativement dans la MÊME "
            "LANGUE que le transcript (pas de traduction).\n"
            "- description : 1 à 2 phrases optimisées SEO TikTok, dans la même langue "
            "que le transcript.\n"
            "- hashtags : 3 à 5 hashtags (mélange niche + large), avec le #.\n"
            "- keywords : 2 à 4 mots forts du clip, recopiés EXACTEMENT tels qu'ils "
            "apparaissent dans le transcript (un seul mot chacun, pas de phrase) — "
            "ils seront surlignés dans les sous-titres.\n"
            "- virality_score : potentiel viral estimé du clip, entier de 0 à 100 "
            "(accroche, émotion, autonomie, partageabilité).\n\n"
            f"Transcript :\n{transcript}"
        )

        client = anthropic.Anthropic()
        response = client.messages.parse(
            model="claude-sonnet-4-6",
            max_tokens=8000,
            thinking={"type": "adaptive"},
            output_config={"effort": "medium"},
            messages=[{"role": "user", "content": prompt}],
            output_format=ClipSelection,
        )
        raw_clips = response.parsed_output.clips
    except Exception as exc:
        _log(f"ℹ️ Sélection IA indisponible ({exc}) → découpage automatique de la vidéo")
        return None

    # Convertir les index de phrases en bornes temporelles et les valider
    n = len(sentences)
    windows: list[dict] = []
    for clip in raw_clips[:max_clips]:
        i = max(0, min(clip.first_sentence, n - 1))
        j = max(i, min(clip.last_sentence, n - 1))
        start = max(0.0, sentences[i]["start"])
        end = min(float(total_duration), sentences[j]["end"])
        if end - start < 10.0 or end - start > 90.0:
            continue
        if any(start < w["end"] and end > w["start"] for w in windows):
            continue  # chevauchement
        windows.append(
            {
                "start": start,
                "end": end,
                "hook": clip.hook.strip(),
                "description": clip.description.strip(),
                "hashtags": [h if h.startswith("#") else f"#{h}" for h in clip.hashtags],
                "keywords": [k.strip() for k in clip.keywords if k.strip()],
                "virality_score": max(0, min(100, clip.virality_score)),
            }
        )

    windows.sort(key=lambda w: w["start"])
    return windows or None