#!/usr/bin/env python3
"""
auto_viral.py
Automatically turns long videos into short viral-style clips.

Features:
- Removes silence
- Finds "high energy" parts (loudness + speech density)
- Punch-in zooms
- Optional subtitles using Whisper
- TikTok/Reels format (1080x1920)
"""

import subprocess, os, librosa, numpy as np, tempfile, json
from moviepy.editor import VideoFileClip, vfx, concatenate_videoclips
from faster_whisper import WhisperModel

# ----------------------------- CONFIG --------------------------------------

TARGET_CLIP_LENGTH = 25        # Seconds, max viral clip length
LOUDNESS_THRESHOLD = -22       # Detect hype moments
MIN_SEGMENT = 1.5              # Minimum segment to keep
ENABLE_SUBTITLES = True
WHISPER_MODEL = "small"

# --------------------------------------------------------------------------

def run(cmd):
    print(">>", cmd)
    subprocess.call(cmd, shell=True)

def extract_audio(video, out_file):
    run(f"ffmpeg -y -i '{video}' -vn -ac 1 -ar 16000 '{out_file}'")

def detect_highlights(audio_file):
    y, sr = librosa.load(audio_file, sr=16000)
    # Extract loudness (RMS)
    frame_length = 2048
    hop = 512
    rms = librosa.feature.rms(y=y, frame_length=frame_length, hop_length=hop)[0]
    times = librosa.frames_to_time(np.arange(len(rms)), sr=sr, hop_length=hop)

    # Threshold high-energy parts
    db = librosa.amplitude_to_db(rms)
    mask = db > LOUDNESS_THRESHOLD

    segments = []
    start = None

    for t, m in zip(times, mask):
        if m and start is None:
            start = t
        elif not m and start:
            if t - start > MIN_SEGMENT:
                segments.append((start, t))
            start = None

    return segments

def choose_best_segment(segments, target_len):
    """Choose the best continuous segment that matches target duration."""
    if not segments:
        return None
    merged = []
    for s,e in segments:
        if not merged:
            merged.append([s,e])
        else:
            if s - merged[-1][1] < 0.5:
                merged[-1][1] = e
            else:
                merged.append([s,e])
    # Find segment closest to target length
    best = max(merged, key=lambda x: min(x[1]-x[0], target_len))
    return best[0], min(best[1], best[0] + target_len)

def generate_subtitles(video, segment_start, segment_end):
    tmp_audio = tempfile.mktemp(suffix=".wav")
    run(f"ffmpeg -y -i '{video}' -ss {segment_start} -t {segment_end-segment_start} -vn '{tmp_audio}'")

    model = WhisperModel(WHISPER_MODEL)
    segments, _ = model.transcribe(tmp_audio)

    srt = "1\n"
    idx = 1
    for seg in segments:
        start = seg.start
        end = seg.end
        text = seg.text.strip()
        srt += f"{idx}\n"
        srt += f"00:00:{start:05.2f} --> 00:00:{end:05.2f}\n"
        srt += text + "\n\n"
        idx += 1

    sub_path = tempfile.mktemp(suffix=".srt")
    open(sub_path,"w").write(srt)
    return sub_path

def edit_video(input_video, out_file):
    # 1) extract audio
    audio_file = tempfile.mktemp(suffix=".wav")
    extract_audio(input_video, audio_file)

    # 2) detect hype segments
    segments = detect_highlights(audio_file)

    if not segments:
        raise SystemExit("No hype detected.")

    # 3) pick best
    start, end = choose_best_segment(segments, TARGET_CLIP_LENGTH)

    # 4) load clip and crop vertical + zoom
    clip = VideoFileClip(input_video).subclip(start, end)

    # center crop to 9:16
    clip = clip.crop(x_center=clip.w/2, y_center=clip.h/2,
                     width=min(clip.w, clip.h*9/16),
                     height=min(clip.h, clip.w*16/9))

    # resize to 1080x1920
    clip = clip.resize((1080,1920))

    # punch-in zoom effect
    clip = clip.fx(vfx.crop, x1=10, y1=10, x2=clip.w-10, y2=clip.h-10)

    # 5) subtitles
    if ENABLE_SUBTITLES:
        sub = generate_subtitles(input_video, start, end)
        tmp = tempfile.mktemp(suffix=".mp4")
        clip.write_videofile(tmp, codec="libx264", audio_codec="aac")

        run(f"ffmpeg -y -i '{tmp}' -vf subtitles='{sub}' -c:a copy '{out_file}'")
    else:
        clip.write_videofile(out_file, codec="libx264", audio_codec="aac")

    print("✔ Exported viral-style clip:", out_file)


if __name__ == "__main__":
    import sys
    if len(sys.argv) < 2:
        print("Usage: python auto_viral.py input.mp4")
        sys.exit(1)

    edit_video(sys.argv[1], "viral_clip.mp4")