import pathlib, subprocess, os, time
from yt_dlp import YoutubeDL
import replicate

MODEL_VER = os.getenv(
    "WHISPER_MODEL",
    "villesau/whisper-timestamped:c5b122b7e513b1b5a6ef849891c538869b77cc932cbd0f8203e11d3b357553b8"
)


# ────────── download TikTok (unchanged) ──────────
def download_tiktok(url: str, out_dir: str) -> str:
    out_dir = pathlib.Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    tmpl = str(out_dir / "%(id)s.%(ext)s")
    ydl_opts = {"outtmpl": tmpl, "format": "mp4/best", "quiet": True, "no_warnings": True}
    with YoutubeDL(ydl_opts) as ydl:
        info = ydl.extract_info(url, download=True)
        return str(pathlib.Path(ydl.prepare_filename(info)).with_suffix(".mp4"))


# ────────── NEW: robust MP4 → PCM‑WAV ──────────
def extract_audio(video_path: str, out_path: str) -> str:
    """
    Extracts **real** 16‑bit PCM WAV at 16 kHz mono – suitable for Whisper & Suno.

    Parameters
    ----------
    video_path : str
        Path to the .mp4 file you just downloaded.
    out_path : str
        Where to write the resulting .wav.

    Returns
    -------
    str
        Same as out_path (handy for piping into the next step).
    """
    cmd = [
        "ffmpeg", "-y",
        "-i", video_path,
        "-ac", "1",          # mono
        "-ar", "16000",      # 16 kHz sample‑rate
        "-c:a", "pcm_s16le", # 16‑bit little‑endian PCM (valid WAV)
        out_path,
    ]
    subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    return out_path


# ────────── transcribe with Replicate Whisper (unchanged) ──────────
def transcribe_with_replicate(clip_id: str) -> dict:
    client = replicate.Client(api_token=os.environ["REPLICATE_API_TOKEN"])
    start = time.time()
    with open(audio_path, "rb") as f:
        out = client.run(
            MODEL_VER,
            input={
                "audio_file": f,
                "language": "auto",
                "task": "transcribe",
                "vad": True,
                "compute_word_confidence": True
            }
        )
    elapsed = time.time() - start
    if isinstance(out, dict):
        out["_elapsed_sec"] = elapsed
    return out