import re
import cv2
import numpy as np
import argparse
import json
from pathlib import Path
from moviepy import VideoFileClip

# ---- CONFIG ----
DEFAULT_KEYWORDS = [
    # Pure Hype - the #1 highlight signal
    "pog", "poggers", "pogchamp", "let's go", "lets go", "huge", "insane", "crazy",
    "unbelievable", "no way", "way too good", "oh my god", "omg", "wow", "holy",
    "holy shit", "holy fuck", "what the fuck", "wtf", "lfg", "clip that", "clip it",
    "someone clip", "chat clip", "world record", "pb", "personal best",

    # Chat interaction = something happened
    "chat", "look at chat", "chat is going crazy", "we are so back", "we're so back",
    "it's over", "ggs", "gg", "not like this", "like this", "AYAYA", "KEKW",
    "monka", "pepega", "copium",

    # Gaming clutch moments
    "clutch", "ace", "1v4", "1v3", "1v2", "quadra", "penta", "pentakill",
    "triple kill", "rampage", "unstoppable", "godlike", "wipe", "team wipe",
    "no scope", "collat", "one shot", "one tap", "200 pumped", "knocked", "cracked",

    # Reaction / laugh / fail
    "bro", "dude", "what is happening", "what just happened", "there is no way",
    "you're kidding", "stop", "no no no", "yes yes yes", "i can't", "i cant believe",
    "this is insane", "this guy", "look at this", "wait wait wait"
]

# You can add your own: --keywords "goal, fight, knock out"
import subprocess
import tempfile
import os
from scipy.io import wavfile

def get_av_scores(video_path, sample_fps=2):
    print(f"-> Opening: {video_path}", flush=True)
    if not os.path.exists(video_path):
        raise FileNotFoundError(video_path)

    with VideoFileClip(video_path) as clip:
        duration = clip.duration
        fps = clip.fps or 30
        has_audio = clip.audio is not None

    print(f" duration={duration:.2f}s fps={fps} has_audio={has_audio}", flush=True)
    n_secs = int(duration)

     # ---- 1. AUDIO (streaming, constant RAM) ----
    print(f"[1/3] Audio energy...", flush=True)
    audio_per_sec = np.zeros(n_secs)

    if has_audio:
        tmp_wav = tempfile.mktemp(suffix=".wav")
        try:
            print(f" extracting mono 16k wav with ffmpeg...", flush=True)
            cmd = [
                "ffmpeg", "-y",
                "-i", video_path,
                "-ar", "16000",
                "-ac", "1",
                "-c:a", "pcm_s16le",
                tmp_wav
            ]
            result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=300)
            if result.returncode!= 0:
                print(f" ffmpeg failed: {result.stderr.decode()[:1000]}")
            else:
                wav_size = os.path.getsize(tmp_wav)
                print(f" wav size: {wav_size} bytes, computing RMS streaming...", flush=True)

                # STREAM instead of loading all at once
                import wave
                with wave.open(tmp_wav, 'rb') as wf:
                    sr = wf.getframerate() # should be 16000
                    n_frames = wf.getnframes()
                    # sr frames = 1 second
                    for i in range(n_secs):
                        if wf.tell() >= n_frames:
                            break
                        # read 1 second
                        frames = wf.readframes(sr)
                        if not frames:
                            break
                        # int16 -> float without extra copy
                        chunk = np.frombuffer(frames, dtype=np.int16).astype(np.float32) / 32768.0
                        if len(chunk) > 0:
                            audio_per_sec[i] = np.sqrt(np.mean(chunk**2))

                        if i % 300 == 0: # log every 5 min
                            print(f" audio {i}/{n_secs}s", flush=True)

                if audio_per_sec.max() > 0:
                    audio_per_sec = (audio_per_sec - audio_per_sec.min()) / (audio_per_sec.max() - audio_per_sec.min() + 1e-6)
                print(f" audio done", flush=True)

        except Exception as e:
            print(f" audio error: {e}")
            import traceback; traceback.print_exc()
        finally:
            if os.path.exists(tmp_wav):
                os.remove(tmp_wav)
    else:
        print(" no audio track", flush=True)

    # ---- 2. MOTION ----
    print(f"[2/3] Motion analysis sample_fps={sample_fps}...", flush=True)
    cap = cv2.VideoCapture(video_path)
    motion_scores = []
    prev_gray = None
    frame_interval = int(fps / sample_fps) if fps > 0 else 15
    num_samples = int(duration * sample_fps)

    for i in range(num_samples):
        cap.set(cv2.CAP_PROP_POS_FRAMES, i * frame_interval)
        ret, frame = cap.read()
        if not ret:
            continue
        gray = cv2.cvtColor(cv2.resize(frame, (320, 180)), cv2.COLOR_BGR2GRAY)
        if prev_gray is not None:
            motion_scores.append(np.mean(cv2.absdiff(gray, prev_gray)) / 255.0)
        else:
            motion_scores.append(0)
        prev_gray = gray
        if i % 50 == 0:
            print(f" motion {i}/{num_samples}", flush=True)
    cap.release()

    if len(motion_scores) > 1:
        x = np.linspace(0, duration-1, len(motion_scores))
        motion_per_sec = np.interp(np.arange(n_secs), x, motion_scores)
        if motion_per_sec.max() > 0:
            motion_per_sec = (motion_per_sec - motion_per_sec.min()) / (motion_per_sec.max() - motion_per_sec.min() + 1e-6)
    else:
        motion_per_sec = np.zeros(n_secs)

    combined_av = 0.6 * audio_per_sec + 0.4 * motion_per_sec
    return combined_av, duration

def srt_time_to_seconds(time_str):
    # 00:01:23,456 or 00:01:23.456
    time_str = time_str.strip().replace('.', ',')
    h, m, s_ms = time_str.split(':')
    s, ms = s_ms.split(',')
    return int(h) * 3600 + int(m) * 60 + int(s) + int(ms) / 1000.0

def load_text_from_srt(srt_path):
    print(f"[3/3] Loading transcript from SRT: {srt_path}")
    if not Path(srt_path).exists():
        raise FileNotFoundError(f"SRT not found: {srt_path}")

    # utf-8-sig strips BOM that Whisper sometimes adds
    with open(srt_path, 'r', encoding='utf-8-sig') as f:
        content = f.read().strip()

    if not content:
        raise ValueError(f"SRT file is empty: {srt_path}")

    # Split by blank line - works for \n\n and \r\n\r\n
    blocks = re.split(r'\r?\n\r?\n', content)
    
    segments = []
    # Pattern inside each block only
    time_pattern = re.compile(r'(\d{2}:\d{2}:\d{2}[,\.]\d{3})\s*-->\s*(\d{2}:\d{2}:\d{2}[,\.]\d{3})')

    for block in blocks:
        block = block.strip()
        if not block:
            continue
        
        lines = block.splitlines()
        if len(lines) < 2:
            continue

        # Find time line (line 0 or 1 - index may be missing)
        time_match = None
        text_start_idx = 0
        for idx, line in enumerate(lines[:3]): # check first 3 lines for time
            m = time_pattern.search(line)
            if m:
                time_match = m
                text_start_idx = idx + 1
                break
        
        if not time_match:
            continue

        start_str, end_str = time_match.groups()
        text = ' '.join(lines[text_start_idx:]).strip()
        
        # Clean Whisper / Twitch artifacts
        text = re.sub(r'<[^>]+>', '', text) # remove <i>, <font> etc
        text = text.replace('\n', ' ').strip()
        
        if text:
            segments.append({
                "start": srt_time_to_seconds(start_str),
                "end": srt_time_to_seconds(end_str),
                "text": text
            })
    
    print(f"Loaded {len(segments)} segments from SRT")
    if len(segments) == 0:
        print(f"WARNING: Parsed 0 segments. First 500 chars of file:\n{content[:500]}")
    return segments

def score_keywords(segments, duration, keywords):
    """Returns per-second keyword boost score + keyword hits"""
    keyword_score = np.zeros(int(duration))
    hits = []

    # Build regex with word boundaries for each keyword
    # \bace\b won't match "place", \bgg\b won't match "ggs"
    patterns = []
    for kw in keywords:
        kw = kw.strip()
        if not kw:
            continue
        # \b at start/end, but handle multi-word phrases
        # re.escape handles "let's go", "1v4", etc
        pat = re.compile(r'\b' + re.escape(kw.lower()) + r'\b', re.IGNORECASE)
        patterns.append((kw.lower(), pat))

    for seg in segments:
        text = seg['text']
        text_lower = text.lower()
        for kw_lower, pat in patterns:
            if pat.search(text_lower):
                start = int(max(0, seg['start'] - 1))
                end = int(min(duration-1, seg['end'] + 2))
                keyword_score[start:end+1] += 1.0
                hits.append({
                    "timestamp": round(seg['start'], 2),
                    "end": round(seg['end'], 2),
                    "keyword": kw_lower,
                    "text": text.strip()
                })
                break

    if keyword_score.max() > 0:
        keyword_score = keyword_score / keyword_score.max()

    return keyword_score, hits

def find_highlights(final_scores, duration, top_n=5, clip_len=8, min_gap=12):
    candidates = []
    for i in range(len(final_scores) - clip_len):
        candidates.append((np.mean(final_scores[i:i+clip_len]), i))

    candidates.sort(reverse=True, key=lambda x: x[0])
    picked = []
    for score, start in candidates:
        if all(abs(start - p[1]) >= min_gap for p in picked):
            picked.append((score, start))
        if len(picked) >= top_n:
            break

    picked.sort(key=lambda x: x[1])
    return [(s, e, sc) for sc, s in [(sc, st) for sc, st in picked] for e in [s+clip_len]]

def format_timestamp(seconds):
    m, s = divmod(int(seconds), 60)
    h, m = divmod(m, 60)
    return f"{h:02d}:{m:02d}:{s:02d}" if h > 0 else f"{m:02d}:{s:02d}"

def highlights(video_path, keywords=None, top_n=5, clip_len=8, min_gap=12):
    if keywords is None:
        keywords = DEFAULT_KEYWORDS

    av_score, duration = get_av_scores(video_path, sample_fps=0.5)
    segments = load_text_from_srt(video_path.replace('.mp4', '.srt')) 
    kw_score, kw_hits = score_keywords(segments, duration, keywords)

    final_score = av_score + (kw_score * 1.5)
    final_score = (final_score - final_score.min()) / (final_score.max() - final_score.min() + 1e-6)
    final_score = np.convolve(final_score, np.ones(3)/3, mode='same')

    highlights = find_highlights(final_score, duration, top_n=top_n, clip_len=clip_len)

    return highlights, kw_hits

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Find highlights with Whisper keywords + timestamps")
    parser.add_argument("video", help="Input video file")
    parser.add_argument("--keywords", type=str, default=",".join(DEFAULT_KEYWORDS), help="Comma-separated keywords to hunt for")
    parser.add_argument("--model", type=str, default="base", help="Whisper model: tiny, base, small, medium")
    parser.add_argument("--top", type=int, default=5)
    parser.add_argument("--len", type=int, default=30, help="Highlight length in seconds")
    parser.add_argument("--no-cut", action="store_true", help="Only print timestamps, don't cut video files")
    args = parser.parse_args()

    keywords = [k.strip() for k in args.keywords.split(",") if k.strip()]

    av_score, duration = get_av_scores(args.video, sample_fps=0.5)
    
    segments = load_text_from_srt(args.video.replace('.mp4', '.srt')) 
    kw_score, kw_hits = score_keywords(segments, duration, keywords)

    # FINAL SCORE: AV + keyword boost (keyword is huge if found)
    final_score = av_score + (kw_score * 1.5)
    final_score = (final_score - final_score.min()) / (final_score.max() - final_score.min() + 1e-6)
    final_score = np.convolve(final_score, np.ones(3)/3, mode='same') # smooth

    highlights = find_highlights(final_score, duration, top_n=args.top, clip_len=args.len)

    # --- OUTPUT ---
    print("\n" + "="*60)
    print("KEYWORD HITS:")
    for h in kw_hits:
        print(f" [{format_timestamp(h['timestamp'])}] ({h['keyword']}) -> {h['text']}")

    print("\n" + "="*60)
    print("FINAL HIGHLIGHTS:")
    output_data = []
    for i, (start, end, score) in enumerate(highlights, 1):
        # find transcript in that window
        texts_in_window = [s['text'].strip() for s in segments if not (s['end'] < start or s['start'] > end)]
        transcript = " ".join(texts_in_window)[:120]

        print(f"{i}. {format_timestamp(start)} - {format_timestamp(end)} (score {score:.2f}) | {transcript}")
        output_data.append({
            "id": i,
            "start_sec": float(start),
            "end_sec": float(end),
            "start_timestamp": format_timestamp(start),
            "end_timestamp": format_timestamp(end),
            "score": float(score),
            "transcript": transcript
        })

    # Save timestamps to JSON
    with open("highlights.json", "w") as f:
        json.dump({"video": args.video, "keywords": keywords, "highlights": output_data, "keyword_hits": kw_hits}, f, indent=2)
    print(f"\nTimestamps saved to highlights.json")

    # Cut videos if requested
    #if not args.no_cut:
    if args.no_cut:
        print("\nCutting clips...")
        Path("highlights").mkdir(exist_ok=True)
        clip = VideoFileClip(args.video)
        for h in output_data:
            sub = clip.subclip(h['start_sec'], h['end_sec'])
            out = f"highlights/highlight_{h['id']}_{h['start_timestamp'].replace(':','-')}.mp4"
            sub.write_videofile(out, codec="libx264", audio_codec="aac", logger=None)
        clip.close()
        print("Done ->./highlights/")