#!/usr/bin/env python3
"""
OPTIMIZED GPU PIPELINE - 100% GPU for AMD with ROCm
Uses: OpenAI Whisper (ROCm) + VAAPI decode/encode
"""
import torch
import json
import subprocess
import os
import re
import whisper

VIDEO = "/home/ren/Documents/proyecto_2026/Twitch-Highlight-Detector/20260219_205705_twitch.mp4"
OUTPUT_DIR = "/home/ren/Documents/proyecto_2026/Twitch-Highlight-Detector"
VADEVICE = "/dev/dri/renderD128"

print("=" * 60)
print("OPTIMIZED GPU PIPELINE - 100% AMD GPU")
print("=" * 60)

if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")

# =====================
# STEP 1: OpenAI Whisper (ROCm GPU)
# =====================
print("\n[1/4] Transcription with Whisper on ROCm GPU...")

transcription_file = f"{OUTPUT_DIR}/transcripcion.json"
if os.path.exists(transcription_file):
    with open(transcription_file) as f:
        transcription = json.load(f)
    print(f"   Loaded {len(transcription.get('segments', []))} segments from cache")
else:
    print("   Loading Whisper small on ROCm GPU...")
    model = whisper.load_model("small", device="cuda")
    print("   Transcribing on GPU...")

    result = model.transcribe(
        VIDEO,
        language="es",
        temperature=0.0
    )

    transcription = {
        'segments': [{'start': s['start'], 'end': s['end'], 'text': s['text']} for s in result['segments']]
    }

    with open(transcription_file, 'w') as f:
        json.dump(transcription, f)

    print(f"   Done: {len(transcription['segments'])} segments")

# =====================
# STEP 2: Analysis (minimal CPU)
# =====================
print("\n[2/4] Analysis...")

patterns = {
    'rage': re.compile(r'put[a-z]*|coño|mierda|hostia|joder|carajo|pendejo', re.I),
    'epic': re.compile(r'kill|matar|muerte|pelea|baron|dragon', re.I),
}

valid = []
for t in transcription.get('segments', []):
    if t.get('start', 0) > 420 and len(t.get('text', '')) > 10:
        text = t.get('text', '').lower()
        if 'dior' not in text:
            score = 0
            for p in patterns.values():
                if p.search(text):
                    score += 15
            if score > 0:
                valid.append({
                    'start': int(t['start']) - 10,
                    'end': int(t['end']) + 20,
                    'score': score
                })

valid.sort(key=lambda x: x['score'], reverse=True)
filtered = []
for v in valid:
    if not filtered or v['start'] - filtered[-1]['end'] > 40:
        filtered.append(v)

filtered = filtered[:20]
print(f"   Highlights: {len(filtered)}")

# =====================
# STEP 3: GPU Encoding (VAAPI decode + encode)
# =====================
print("\n[3/4] GPU Encoding (VAAPI decode + encode)...")

def encode_clip(i, h):
    output = f"{OUTPUT_DIR}/opt_clip_{i}.mp4"
    start = h['start']
    duration = min(h['end'] - h['start'], 60)

    cmd = [
        'ffmpeg', '-y',
        '-hwaccel', 'vaapi',
        '-hwaccel_device', VADEVICE,
        '-hwaccel_output_format', 'vaapi',
        '-i', VIDEO,
        '-ss', str(start),
        '-t', str(duration),
        '-vf', 'scale_vaapi=1920:1080:format=nv12',
        '-c:v', 'h264_vaapi',
        '-b:v', '5M',
        '-c:a', 'aac', '-b:a', '128k',
        output
    ]
    
    result = subprocess.run(cmd, capture_output=True)
    if os.path.exists(output) and os.path.getsize(output) > 1000:
        print(f"   Clip {i} done (VAAPI GPU)")
        return output
    
    stderr = result.stderr.decode()[-300:] if result.stderr else "no error"
    print(f"   Clip {i} VAAPI failed: {stderr}")
    return None

clips = []
for i, h in enumerate(filtered):
    result = encode_clip(i, h)
    if result:
        clips.append(result)

print(f"   Encoded {len(clips)} clips on GPU")

# =====================
# STEP 4: Concatenate
# =====================
print("\n[4/4] Concatenating...")

if not clips:
    print("   No clips to concatenate!")
    exit(1)

concat_file = f"{OUTPUT_DIR}/concat_opt.txt"
with open(concat_file, 'w') as f:
    for clip in clips:
        f.write(f"file '{clip}'\n")

output = f"{OUTPUT_DIR}/HIGHLIGHTS_OPT.mp4"

cmd = [
    'ffmpeg', '-y',
    '-f', 'concat', '-safe', '0', '-i', concat_file,
    '-c', 'copy',
    output
]
subprocess.run(cmd, capture_output=True)

# Cleanup
for c in clips:
    if os.path.exists(c):
        os.remove(c)
if os.path.exists(concat_file):
    os.remove(concat_file)

if os.path.exists(output):
    size = os.path.getsize(output) / 1024 / 1024
    print(f"\n✅ DONE! {output} ({size:.1f} MB)")
else:
    print("ERROR creating final file!")