import os
# Force multi-threading in math/ML libraries to use all CPU resources
os.environ["OMP_NUM_THREADS"] = "8"
os.environ["MKL_NUM_THREADS"] = "8"
os.environ["OPENBLAS_NUM_THREADS"] = "8"
os.environ["VECLIB_MAXIMUM_THREADS"] = "8"
os.environ["NUMEXPR_NUM_THREADS"] = "8"

import time
import cv2
import scenedetect
import subprocess
import argparse
import re
import sys
from scenedetect import open_video, SceneManager
from scenedetect.detectors import ContentDetector
from ultralytics import YOLO
import torch
import os
import numpy as np
from tqdm import tqdm
import yt_dlp
import mediapipe as mp
# import whisper (replaced by faster_whisper inside function)
from google import genai
from dotenv import load_dotenv
import json
from subtitles import generate_srt, burn_subtitles

import warnings
warnings.filterwarnings("ignore", category=UserWarning, module='google.protobuf')

# Load environment variables
load_dotenv()

# --- Constants ---
ASPECT_RATIO = 9 / 16

GEMINI_PROMPT_TEMPLATE = """
You are a senior short-form video editor. Read the ENTIRE transcript and word-level timestamps to choose the 3–15 MOST VIRAL moments for TikTok/IG Reels/YouTube Shorts. Each clip must be between 15 and 60 seconds long.

⚠️ FFMPEG TIME CONTRACT — STRICT REQUIREMENTS:
- Return timestamps in ABSOLUTE SECONDS from the start of the video (usable in: ffmpeg -ss <start> -to <end> -i <input> ...).
- Only NUMBERS with decimal point, up to 3 decimals (examples: 0, 1.250, 17.350).
- Ensure 0 ≤ start < end ≤ VIDEO_DURATION_SECONDS.
- Each clip between 15 and 60 s (inclusive).
- Prefer starting 0.2–0.4 s BEFORE the hook and ending 0.2–0.4 s AFTER the payoff.
- Use silence moments for natural cuts; never cut in the middle of a word or phrase.
- STRICTLY FORBIDDEN to use time formats other than absolute seconds.

VIDEO_DURATION_SECONDS: {video_duration}

TRANSCRIPT_TEXT (raw):
{transcript_text}

WORDS_JSON (array of {{w, s, e}} where s/e are seconds):
{words_json}

STRICT EXCLUSIONS:
- No generic intros/outros or purely sponsorship segments unless they contain the hook.
- No clips < 15 s or > 60 s.

OUTPUT — RETURN ONLY VALID JSON (no markdown, no comments). Order clips by predicted performance (best to worst). In the descriptions, ALWAYS include a CTA like "Follow me and comment X and I'll send you the workflow" (especially if discussing an n8n workflow):
{{
  "shorts": [
    {{
      "start": <number in seconds, e.g., 12.340>,
      "end": <number in seconds, e.g., 37.900>,
      "video_description_for_tiktok": "<description for TikTok oriented to get views>",
      "video_description_for_instagram": "<description for Instagram oriented to get views>",
      "video_title_for_youtube_short": "<title for YouTube Short oriented to get views 100 chars max>",
      "viral_hook_text": "<SHORT punchy text overlay (max 10 words). MUST BE IN THE SAME LANGUAGE AS THE VIDEO TRANSCRIPT. Examples: 'POV: You realized...', 'Did you know?', 'Stop doing this!'>"
    }}
  ]
}}
"""

import threading

# --- MediaPipe Setup ---
mp_face_detection = mp.solutions.face_detection

thread_local = threading.local()

def get_face_detector():
    if not hasattr(thread_local, 'detector'):
        thread_local.detector = mp_face_detection.FaceDetection(model_selection=1, min_detection_confidence=0.5)
    return thread_local.detector

def get_yolo_model():
    if not hasattr(thread_local, 'yolo'):
        thread_local.yolo = YOLO('yolov8n.pt')
    return thread_local.yolo

class SmoothedCameraman:
    """
    Handles smooth camera movement.
    Simplified Logic: "Heavy Tripod"
    Only moves if the subject leaves the center safe zone.
    Moves slowly and linearly.
    """
    def __init__(self, output_width, output_height, video_width, video_height):
        self.output_width = output_width
        self.output_height = output_height
        self.video_width = video_width
        self.video_height = video_height
        
        # Initial State
        self.current_center_x = video_width / 2
        self.target_center_x = video_width / 2
        
        # Calculate crop dimensions once
        self.crop_height = video_height
        self.crop_width = int(self.crop_height * ASPECT_RATIO)
        if self.crop_width > video_width:
             self.crop_width = video_width
             self.crop_height = int(self.crop_width / ASPECT_RATIO)
             
        # Safe Zone: 20% of the video width
        # As long as the target is within this zone relative to current center, DO NOT MOVE.
        self.safe_zone_radius = self.crop_width * 0.25

    def update_target(self, face_box):
        """
        Updates the target center based on detected face/person.
        """
        if face_box:
            x, y, w, h = face_box
            self.target_center_x = x + w / 2
    
    def get_crop_box(self, force_snap=False):
        """
        Returns the (x1, y1, x2, y2) for the current frame.
        """
        if force_snap:
            self.current_center_x = self.target_center_x
        else:
            diff = self.target_center_x - self.current_center_x
            
            # SIMPLIFIED LOGIC:
            # 1. Is the target outside the safe zone?
            if abs(diff) > self.safe_zone_radius:
                # 2. If yes, move towards it slowly (Linear Speed)
                # Determine direction
                direction = 1 if diff > 0 else -1
                
                # Speed: 2 pixels per frame (Slow pan)
                # If the distance is HUGE (scene change or fast movement), speed up slightly
                if abs(diff) > self.crop_width * 0.5:
                    speed = 15.0 # Fast re-frame
                else:
                    speed = 3.0  # Slow, steady pan
                
                self.current_center_x += direction * speed
                
                # Check if we overshot (prevent oscillation)
                new_diff = self.target_center_x - self.current_center_x
                if (direction == 1 and new_diff < 0) or (direction == -1 and new_diff > 0):
                    self.current_center_x = self.target_center_x
            
            # If inside safe zone, DO NOTHING (Stationary Camera)
                
        # Clamp center
        half_crop = self.crop_width / 2
        
        if self.current_center_x - half_crop < 0:
            self.current_center_x = half_crop
        if self.current_center_x + half_crop > self.video_width:
            self.current_center_x = self.video_width - half_crop
            
        x1 = int(self.current_center_x - half_crop)
        x2 = int(self.current_center_x + half_crop)
        
        x1 = max(0, x1)
        x2 = min(self.video_width, x2)
        
        y1 = 0
        y2 = self.video_height
        
        return x1, y1, x2, y2

class SpeakerTracker:
    """
    Tracks speakers over time to prevent rapid switching and handle temporary obstructions.
    """
    def __init__(self, stabilization_frames=15, cooldown_frames=30):
        self.active_speaker_id = None
        self.speaker_scores = {}  # {id: score}
        self.last_seen = {}       # {id: frame_number}
        self.locked_counter = 0   # How long we've been locked on current speaker
        
        # Hyperparameters
        self.stabilization_threshold = stabilization_frames # Frames needed to confirm a new speaker
        self.switch_cooldown = cooldown_frames              # Minimum frames before switching again
        self.last_switch_frame = -1000
        
        # ID tracking
        self.next_id = 0
        self.known_faces = [] # [{'id': 0, 'center': x, 'last_frame': 123}]

    def get_target(self, face_candidates, frame_number, width):
        """
        Decides which face to focus on.
        face_candidates: list of {'box': [x,y,w,h], 'score': float}
        """
        current_candidates = []
        
        # 1. Match faces to known IDs (simple distance tracking)
        for face in face_candidates:
            x, y, w, h = face['box']
            center_x = x + w / 2
            
            best_match_id = -1
            min_dist = width * 0.15 # Reduced matching radius to avoid jumping in groups
            
            # Try to match with known faces seen recently
            for kf in self.known_faces:
                if frame_number - kf['last_frame'] > 30: # Forgot faces older than 1s (was 2s)
                    continue
                    
                dist = abs(center_x - kf['center'])
                if dist < min_dist:
                    min_dist = dist
                    best_match_id = kf['id']
            
            # If no match, assign new ID
            if best_match_id == -1:
                best_match_id = self.next_id
                self.next_id += 1
            
            # Update known face
            self.known_faces = [kf for kf in self.known_faces if kf['id'] != best_match_id]
            self.known_faces.append({'id': best_match_id, 'center': center_x, 'last_frame': frame_number})
            
            current_candidates.append({
                'id': best_match_id,
                'box': face['box'],
                'score': face['score']
            })

        # 2. Update Scores with decay
        for pid in list(self.speaker_scores.keys()):
             self.speaker_scores[pid] *= 0.85 # Faster decay (was 0.9)
             if self.speaker_scores[pid] < 0.1:
                 del self.speaker_scores[pid]

        # Add new scores
        for cand in current_candidates:
            pid = cand['id']
            # Score is purely based on size (proximity) now that we don't have mouth
            raw_score = cand['score'] / (width * width * 0.05)
            self.speaker_scores[pid] = self.speaker_scores.get(pid, 0) + raw_score

        # 3. Determine Best Speaker
        if not current_candidates:
            # If no one found, maintain last active speaker if cooldown allows
            # to avoid black screen or jump to 0,0
            return None 
            
        best_candidate = None
        max_score = -1
        
        for cand in current_candidates:
            pid = cand['id']
            total_score = self.speaker_scores.get(pid, 0)
            
            # Hysteresis: HUGE Bonus for current active speaker
            if pid == self.active_speaker_id:
                total_score *= 3.0 # Sticky factor
                
            if total_score > max_score:
                max_score = total_score
                best_candidate = cand

        # 4. Decide Switch
        if best_candidate:
            target_id = best_candidate['id']
            
            if target_id == self.active_speaker_id:
                self.locked_counter += 1
                return best_candidate['box']
            
            # New person
            if frame_number - self.last_switch_frame < self.switch_cooldown:
                old_cand = next((c for c in current_candidates if c['id'] == self.active_speaker_id), None)
                if old_cand:
                    return old_cand['box']
            
            self.active_speaker_id = target_id
            self.last_switch_frame = frame_number
            self.locked_counter = 0
            return best_candidate['box']
            
        return None

def detect_face_candidates(frame):
    """
    Returns list of all detected faces using lightweight FaceDetection.
    """
    height, width, _ = frame.shape
    rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    detector = get_face_detector()
    results = detector.process(rgb_frame)
    
    candidates = []
    
    if not results.detections:
        return []
        
    for detection in results.detections:
        bboxC = detection.location_data.relative_bounding_box
        x = int(bboxC.xmin * width)
        y = int(bboxC.ymin * height)
        w = int(bboxC.width * width)
        h = int(bboxC.height * height)
        
        candidates.append({
            'box': [x, y, w, h],
            'score': w * h # Area as score
        })
            
    return candidates

def detect_person_yolo(frame):
    """
    Fallback: Detect largest person using YOLO when face detection fails.
    Returns [x, y, w, h] of the person's 'upper body' approximation.
    """
    yolo_model = get_yolo_model()
    results = yolo_model(frame, verbose=False, classes=[0]) # class 0 is person
    
    if not results:
        return None
        
    best_box = None
    max_area = 0
    
    for result in results:
        boxes = result.boxes
        for box in boxes:
            x1, y1, x2, y2 = [int(i) for i in box.xyxy[0]]
            w = x2 - x1
            h = y2 - y1
            area = w * h
            
            if area > max_area:
                max_area = area
                # Focus on the top 40% of the person (head/chest) for framing
                # This approximates where the face is if we can't detect it directly
                face_h = int(h * 0.4)
                best_box = [x1, y1, w, face_h]
                
    return best_box

def create_general_frame(frame, output_width, output_height):
    """
    Creates a 'General Shot' frame: 
    - Background: Blurred zoom of original
    - Foreground: Original video scaled to fit width, centered vertically.
    """
    orig_h, orig_w = frame.shape[:2]
    
    # 1. Background (Fill Height)
    # Crop center to aspect ratio
    bg_scale = output_height / orig_h
    bg_w = int(orig_w * bg_scale)
    bg_resized = cv2.resize(frame, (bg_w, output_height))
    
    # Crop center of background
    start_x = (bg_w - output_width) // 2
    if start_x < 0: start_x = 0
    background = bg_resized[:, start_x:start_x+output_width]
    if background.shape[1] != output_width:
        background = cv2.resize(background, (output_width, output_height))
        
    # Blur background
    background = cv2.GaussianBlur(background, (51, 51), 0)
    
    # 2. Foreground (Fit Width)
    scale = output_width / orig_w
    fg_h = int(orig_h * scale)
    foreground = cv2.resize(frame, (output_width, fg_h))
    
    # 3. Overlay
    y_offset = (output_height - fg_h) // 2
    
    # Clone background to avoid modifying it
    final_frame = background.copy()
    final_frame[y_offset:y_offset+fg_h, :] = foreground
    
    return final_frame

def analyze_scenes_strategy(video_path, scenes):
    """
    Analyzes each scene to determine if it should be TRACK (Single person) or GENERAL (Group/Wide).
    Returns list of strategies corresponding to scenes.
    """
    cap = cv2.VideoCapture(video_path)
    strategies = []
    
    if not cap.isOpened():
        return ['TRACK'] * len(scenes)
        
    for start, end in tqdm(scenes, desc="   Analyzing Scenes"):
        # Sample 3 frames (start, middle, end)
        frames_to_check = [
            start.get_frames() + 5,
            int((start.get_frames() + end.get_frames()) / 2),
            end.get_frames() - 5
        ]
        
        face_counts = []
        for f_idx in frames_to_check:
            cap.set(cv2.CAP_PROP_POS_FRAMES, f_idx)
            ret, frame = cap.read()
            if not ret: continue
            
            # Detect faces
            candidates = detect_face_candidates(frame)
            face_counts.append(len(candidates))
            
        # Decision Logic
        if not face_counts:
            avg_faces = 0
        else:
            avg_faces = sum(face_counts) / len(face_counts)
            
        # Strategy:
        # 0 faces -> GENERAL (Landscape/B-roll)
        # 1 face -> TRACK
        # > 1.2 faces -> GENERAL (Group)
        
        if avg_faces > 1.2 or avg_faces < 0.5:
            strategies.append('GENERAL')
        else:
            strategies.append('TRACK')
            
    cap.release()
    return strategies

def detect_scenes(video_path):
    video = open_video(video_path)
    scene_manager = SceneManager()
    scene_manager.add_detector(ContentDetector())
    scene_manager.detect_scenes(video=video)
    scene_list = scene_manager.get_scene_list()
    fps = video.frame_rate
    return scene_list, fps

def get_video_resolution(video_path):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise IOError(f"Could not open video file {video_path}")
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    cap.release()
    return width, height

def sanitize_filename(filename):
    """Remove invalid characters from filename."""
    filename = re.sub(r'[<>:"/\\|?*#]', '', filename)
    filename = filename.replace(' ', '_')
    return filename[:100]


def download_youtube_video(url, output_dir="."):
    """
    Downloads a YouTube video using the custom fixed yt-dl wrapper in /home/ren/yt/yt-dl.
    """
    import subprocess
    import json
    import time
    import sys
    
    print("📥 Downloading video using custom solver...")
    step_start_time = time.time()
    
    # Ensure output_dir is absolute path
    abs_output_dir = os.path.abspath(output_dir)
    
    # We run the command with --json to get structured output
    cmd = [
        "/home/ren/yt/venv/bin/python",
        "/home/ren/yt/yt-dl",
        "--json",
        "-o", abs_output_dir,
        url
    ]
    
    print(f"   Executing: {' '.join(cmd)}")
    result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    
    if result.returncode != 0:
        print("🚨 DOWNLOAD ERROR 🚨", file=sys.stderr)
        print(result.stderr, file=sys.stderr)
        raise Exception(f"Video download failed: {result.stderr}")
        
    try:
        output_data = json.loads(result.stdout.strip())
        downloaded_file = output_data['filepath']
        video_title = output_data.get('title', 'video')
        sanitized_title = sanitize_filename(video_title)
    except Exception as e:
        print(f"⚠️ JSON parsing failed, trying fallback detection. Error: {e}")
        print(f"   Raw stdout: {result.stdout}")
        
        # Fallback parsing: if not JSON, the script prints the absolute filepath to stdout
        downloaded_file = result.stdout.strip()
        if not os.path.exists(downloaded_file):
            # Search output_dir for files created recently
            downloaded_file = None
            for file in os.listdir(abs_output_dir):
                if file.endswith('.mp4'):
                    full_path = os.path.join(abs_output_dir, file)
                    if time.time() - os.path.getmtime(full_path) < 60:
                        downloaded_file = full_path
                        break
                        
        if not downloaded_file or not os.path.exists(downloaded_file):
            raise Exception("Failed to locate downloaded video file.")
            
        sanitized_title = sanitize_filename(os.path.splitext(os.path.basename(downloaded_file))[0])
        
    step_end_time = time.time()
    print(f"✅ Video downloaded in {step_end_time - step_start_time:.2f}s: {downloaded_file}")
    return downloaded_file, sanitized_title

def process_video_to_vertical(input_video, final_output_video, mute_ranges=None):
    """
    Core logic to convert horizontal video to vertical using scene detection and Active Speaker Tracking (MediaPipe).
    """
    script_start_time = time.time()
    
    # Define temporary file paths based on the output name
    base_name = os.path.splitext(final_output_video)[0]
    temp_video_output = f"{base_name}_temp_video.mp4"
    temp_audio_output = f"{base_name}_temp_audio.aac"
    
    # Clean up previous temp files if they exist
    if os.path.exists(temp_video_output): os.remove(temp_video_output)
    if os.path.exists(temp_audio_output): os.remove(temp_audio_output)
    if os.path.exists(final_output_video): os.remove(final_output_video)

    print(f"🎬 Processing clip: {input_video}")
    print("   Step 1: Detecting scenes...")
    scenes, fps = detect_scenes(input_video)
    
    if not scenes:
        print("   ❌ No scenes were detected. Using full video as one scene.")
        # If scene detection fails or finds nothing, treat whole video as one scene
        cap = cv2.VideoCapture(input_video)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        cap.release()
        from scenedetect import FrameTimecode
        scenes = [(FrameTimecode(0, fps), FrameTimecode(total_frames, fps))]

    print(f"   ✅ Found {len(scenes)} scenes.")

    print("\n   🧠 Step 2: Preparing Active Tracking...")
    original_width, original_height = get_video_resolution(input_video)
    
    OUTPUT_HEIGHT = original_height
    OUTPUT_WIDTH = int(OUTPUT_HEIGHT * ASPECT_RATIO)
    if OUTPUT_WIDTH % 2 != 0:
        OUTPUT_WIDTH += 1

    # Initialize Cameraman
    cameraman = SmoothedCameraman(OUTPUT_WIDTH, OUTPUT_HEIGHT, original_width, original_height)
    
    # --- New Strategy: Per-Scene Analysis ---
    print("\n   🤖 Step 3: Analyzing Scenes for Strategy (Single vs Group)...")
    scene_strategies = analyze_scenes_strategy(input_video, scenes)
    # scene_strategies is a list of 'TRACK' or 'General' corresponding to scenes
    
    print("\n   ✂️ Step 4: Processing video frames...")
    
    command = [
        'ffmpeg', '-y', '-f', 'rawvideo', '-vcodec', 'rawvideo',
        '-s', f'{OUTPUT_WIDTH}x{OUTPUT_HEIGHT}', '-pix_fmt', 'bgr24',
        '-r', str(fps), '-i', '-', '-c:v', 'libx264',
        '-preset', 'fast', '-crf', '23', '-an', temp_video_output
    ]

    ffmpeg_process = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)

    cap = cv2.VideoCapture(input_video)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    
    frame_number = 0
    current_scene_index = 0
    
    # Pre-calculate scene boundaries
    scene_boundaries = []
    for s_start, s_end in scenes:
        scene_boundaries.append((s_start.get_frames(), s_end.get_frames()))

    # Global tracker for single-person shots
    speaker_tracker = SpeakerTracker(cooldown_frames=30)

    with tqdm(total=total_frames, desc="   Processing", file=sys.stdout) as pbar:
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            # Update Scene Index
            if current_scene_index < len(scene_boundaries):
                start_f, end_f = scene_boundaries[current_scene_index]
                if frame_number >= end_f and current_scene_index < len(scene_boundaries) - 1:
                    current_scene_index += 1
            
            # Determine Strategy for current frame based on scene
            current_strategy = scene_strategies[current_scene_index] if current_scene_index < len(scene_strategies) else 'TRACK'
            
            # Apply Strategy
            if current_strategy == 'GENERAL':
                # "Plano General" -> Blur Background + Fit Width
                output_frame = create_general_frame(frame, OUTPUT_WIDTH, OUTPUT_HEIGHT)
                
                # Reset cameraman/tracker so they don't drift while inactive
                cameraman.current_center_x = original_width / 2
                cameraman.target_center_x = original_width / 2
                
            else:
                # "Single Speaker" -> Track & Crop
                
                # Detect every 2nd frame for performance
                if frame_number % 2 == 0:
                    candidates = detect_face_candidates(frame)
                    target_box = speaker_tracker.get_target(candidates, frame_number, original_width)
                    if target_box:
                        cameraman.update_target(target_box)
                    else:
                        person_box = detect_person_yolo(frame)
                        if person_box:
                            cameraman.update_target(person_box)

                # Snap camera on scene change to avoid panning from previous scene position
                is_scene_start = (frame_number == scene_boundaries[current_scene_index][0])
                
                x1, y1, x2, y2 = cameraman.get_crop_box(force_snap=is_scene_start)
                
                # Crop
                if y2 > y1 and x2 > x1:
                    cropped = frame[y1:y2, x1:x2]
                    output_frame = cv2.resize(cropped, (OUTPUT_WIDTH, OUTPUT_HEIGHT))
                else:
                    output_frame = cv2.resize(frame, (OUTPUT_WIDTH, OUTPUT_HEIGHT))

            ffmpeg_process.stdin.write(output_frame.tobytes())
            frame_number += 1
            pbar.update(1)
    
    ffmpeg_process.stdin.close()
    stderr_output = ffmpeg_process.stderr.read().decode()
    ffmpeg_process.wait()
    cap.release()

    if ffmpeg_process.returncode != 0:
        print("\n   ❌ FFmpeg frame processing failed.")
        print("   Stderr:", stderr_output)
        return False

    print("\n   🔊 Step 5: Extracting audio...")
    if mute_ranges:
        af_parts = []
        for ms_start, ms_end in mute_ranges:
            af_parts.append(
                f"volume=enable='between(t,{ms_start:.3f},{ms_end:.3f})':volume=0"
            )
        af_filter = ",".join(af_parts)
        audio_extract_command = [
            'ffmpeg', '-y', '-i', input_video, '-vn', '-af', af_filter, '-c:a', 'aac', temp_audio_output
        ]
    else:
        audio_extract_command = [
            'ffmpeg', '-y', '-i', input_video, '-vn', '-c:a', 'aac', temp_audio_output
        ]
        
    try:
        subprocess.run(audio_extract_command, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
    except subprocess.CalledProcessError:
        print("\n   ❌ Audio extraction failed (maybe no audio?). Proceeding without audio.")
        pass

    print("\n   ✨ Step 6: Merging...")
    if os.path.exists(temp_audio_output):
        merge_command = [
            'ffmpeg', '-y', '-i', temp_video_output, '-i', temp_audio_output,
            '-c:v', 'copy', '-c:a', 'copy', final_output_video
        ]
    else:
         merge_command = [
            'ffmpeg', '-y', '-i', temp_video_output,
            '-c:v', 'copy', final_output_video
        ]
        
    try:
        subprocess.run(merge_command, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
        print(f"   ✅ Clip saved to {final_output_video}")
    except subprocess.CalledProcessError as e:
        print("\n   ❌ Final merge failed.")
        print("   Stderr:", e.stderr.decode())
        return False

    # Clean up temp files
    if os.path.exists(temp_video_output): os.remove(temp_video_output)
    if os.path.exists(temp_audio_output): os.remove(temp_audio_output)
    
    return True

def transcribe_chunk(chunk_info):
    video_path, start_time, duration, chunk_index, output_dir = chunk_info
    chunk_audio_path = os.path.join(output_dir, f"chunk_{chunk_index}.wav")
    
    # Extract audio chunk
    cmd = [
        'ffmpeg', '-y',
        '-ss', str(start_time),
        '-t', str(duration),
        '-i', video_path,
        '-vn', '-acodec', 'pcm_s16le', '-ar', '16000', '-ac', '1',
        chunk_audio_path
    ]
    subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
    
    if not os.path.exists(chunk_audio_path):
        return []
        
    from faster_whisper import WhisperModel
    # Each chunk worker uses 2 threads
    model = WhisperModel("base", device="cpu", compute_type="int8", cpu_threads=2)
    segments, info = model.transcribe(chunk_audio_path, word_timestamps=True)
    
    chunk_segments = []
    for segment in segments:
        seg_dict = {
            'text': segment.text,
            'start': segment.start + start_time,
            'end': segment.end + start_time,
            'words': []
        }
        if segment.words:
            for word in segment.words:
                seg_dict['words'].append({
                    'word': word.word,
                    'start': word.start + start_time,
                    'end': word.end + start_time,
                    'probability': word.probability
                })
        chunk_segments.append(seg_dict)
        
    # Clean up chunk file
    if os.path.exists(chunk_audio_path):
        os.remove(chunk_audio_path)
        
    return chunk_segments

def transcribe_video(video_path):
    print("🎙️  Transcribing video with Faster-Whisper in parallel (CPU Optimized)...")
    import cv2
    from concurrent.futures import ThreadPoolExecutor
    
    # 1. Get video duration
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    duration = frame_count / fps if fps else 0
    cap.release()
    
    if duration == 0:
        raise ValueError(f"Could not read duration of {video_path}")
        
    # 2. Divide into 5-minute chunks
    chunk_size = 300.0
    chunks = []
    start = 0.0
    chunk_index = 0
    output_dir = os.path.dirname(video_path) or "."
    
    while start < duration:
        length = min(chunk_size, duration - start)
        chunks.append((video_path, start, length, chunk_index, output_dir))
        start += chunk_size
        chunk_index += 1
        
    print(f"   Splitting video into {len(chunks)} chunks for parallel transcription to use all CPU cores...")
    
    all_segments = []
    # Use ThreadPoolExecutor to run all chunk transcribing concurrently
    with ThreadPoolExecutor(max_workers=len(chunks)) as executor:
        results = list(executor.map(transcribe_chunk, chunks))
        
    for chunk_segs in results:
        all_segments.extend(chunk_segs)
        
    # Sort segments by start time
    all_segments = sorted(all_segments, key=lambda x: x['start'])
    
    # Print the segments to the log
    for seg in all_segments:
        print(f"   [{seg['start']:.2f}s -> {seg['end']:.2f}s] {seg['text']}")
        
    # Reconstruct full text
    full_text = " ".join(seg['text'] for seg in all_segments)
    
    print(f"✅ Parallel transcription complete! Total segments: {len(all_segments)}")
    return {
        'text': full_text.strip(),
        'segments': all_segments,
        'language': 'es'
    }

def get_viral_clips(transcript_result, video_duration):
    print("🤖 Calling local LLM (deepseek-v4-flash-free via opencode-proxy)...")
    import json
    
    # Extract words
    words = []
    for segment in transcript_result['segments']:
        for word in segment.get('words', []):
            words.append({
                'w': word['word'],
                's': word['start'],
                'e': word['end']
            })

    prompt = GEMINI_PROMPT_TEMPLATE.format(
        video_duration=video_duration,
        transcript_text=json.dumps(transcript_result['text']),
        words_json=json.dumps(words)
    )

    from openai import OpenAI
    client = OpenAI(
        base_url="http://localhost:6446/v1",
        api_key="1bf4348a7b1116aeccdcf2583c5390f59ee7eab7ea40dcdc"
    )
    
    # Try with JSON mode first
    text = ""
    try:
        response = client.chat.completions.create(
            model="deepseek-v4-flash-free",
            messages=[{"role": "user", "content": prompt}],
            response_format={"type": "json_object"}
        )
        text = response.choices[0].message.content
    except Exception as e:
        print(f"⚠️ JSON Mode failed or unsupported: {e}. Retrying without JSON mode...")
        try:
            response = client.chat.completions.create(
                model="deepseek-v4-flash-free",
                messages=[{"role": "user", "content": prompt}]
            )
            text = response.choices[0].message.content
        except Exception as e2:
            print(f"❌ Local LLM call failed completely: {e2}")
            return None

    # Clean response if it contains markdown code blocks
    if text.startswith("```json"):
        text = text[7:]
    if text.startswith("```"):
        text = text[3:]
    if text.endswith("```"):
        text = text[:-3]
    text = text.strip()
    
    try:
        result_json = json.loads(text)
        return result_json
    except Exception as e:
        print(f"❌ Failed to parse JSON from LLM response: {e}")
        print("Raw response was:")
        print(text)
        return None

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="AutoCrop-Vertical with Viral Clip Detection.")
    
    input_group = parser.add_mutually_exclusive_group(required=True)
    input_group.add_argument('-i', '--input', type=str, help="Path to the input video file.")
    input_group.add_argument('-u', '--url', type=str, help="YouTube URL to download and process.")
    
    parser.add_argument('-o', '--output', type=str, help="Output directory or file (if processing whole video).")
    parser.add_argument('--keep-original', action='store_true', help="Keep the downloaded YouTube video.")
    parser.add_argument('--skip-analysis', action='store_true', help="Skip AI analysis and convert the whole video.")
    parser.add_argument('--speed', type=float, default=1.25, help="Speed up factor for the output clips (default: 1.25).")
    
    args = parser.parse_args()

    script_start_time = time.time()
    
    def _ensure_dir(path: str) -> str:
        """Create directory if missing and return the same path."""
        if path:
            os.makedirs(path, exist_ok=True)
        return path
    
    # 1. Get Input Video
    if args.url:
        # For multi-clip runs, treat --output as an OUTPUT DIRECTORY (create it if needed).
        # For whole-video runs (--skip-analysis), --output can be a file path.
        if args.output and not args.skip_analysis:
            output_dir = _ensure_dir(args.output)
        else:
            # If output is a directory, use it; if it's a filename, use its directory; else default "."
            if args.output and os.path.isdir(args.output):
                output_dir = args.output
            elif args.output and not os.path.isdir(args.output):
                output_dir = os.path.dirname(args.output) or "."
            else:
                output_dir = "."
        
        input_video, video_title = download_youtube_video(args.url, output_dir)
    else:
        input_video = args.input
        video_title = os.path.splitext(os.path.basename(input_video))[0]
        
        if args.output and not args.skip_analysis:
            # For multi-clip runs, treat --output as an OUTPUT DIRECTORY (create it if needed).
            output_dir = _ensure_dir(args.output)
        else:
            # If output is a directory, use it; if it's a filename, use its directory; else default to input dir.
            if args.output and os.path.isdir(args.output):
                output_dir = args.output
            elif args.output and not os.path.isdir(args.output):
                output_dir = os.path.dirname(args.output) or os.path.dirname(input_video)
            else:
                output_dir = os.path.dirname(input_video)

    if not os.path.exists(input_video):
        print(f"❌ Input file not found: {input_video}")
        exit(1)

    # 2. Decision: Analyze clips or process whole?
    if args.skip_analysis:
        print("⏩ Skipping analysis, processing entire video...")
        output_file = args.output if args.output else os.path.join(output_dir, f"{video_title}_vertical.mp4")
        process_video_to_vertical(input_video, output_file)
    else:
        metadata_file = os.path.join(output_dir, f"{video_title}_metadata.json")
        clips_data = None
        transcript = None
        
        if os.path.exists(metadata_file):
            print(f"📦 Found existing metadata file: {metadata_file}. Loading...")
            try:
                with open(metadata_file, 'r', encoding='utf-8') as f:
                    clips_data = json.load(f)
                if clips_data and 'shorts' in clips_data:
                    transcript = clips_data.get('transcript')
                    print(f"   Successfully loaded {len(clips_data['shorts'])} clips from metadata.")
                else:
                    clips_data = None
            except Exception as e:
                print(f"   ⚠️ Failed to load metadata file: {e}")
                clips_data = None
                
        if not clips_data:
            # 3. Transcribe
            transcript = transcribe_video(input_video)
            
            # Get duration
            cap = cv2.VideoCapture(input_video)
            fps = cap.get(cv2.CAP_PROP_FPS)
            frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            duration = frame_count / fps
            cap.release()

            # 4. Gemini/OpenAI Analysis
            clips_data = get_viral_clips(transcript, duration)
            
        if not clips_data or 'shorts' not in clips_data:
            print("❌ Failed to identify clips. Converting whole video as fallback.")
            output_file = os.path.join(output_dir, f"{video_title}_vertical.mp4")
            process_video_to_vertical(input_video, output_file)
        else:
            print(f"🔥 Found {len(clips_data['shorts'])} viral clips!")
            
            # Save metadata
            clips_data['transcript'] = transcript # Save full transcript for subtitles
            metadata_file = os.path.join(output_dir, f"{video_title}_metadata.json")
            with open(metadata_file, 'w') as f:
                json.dump(clips_data, f, indent=2)
            print(f"   Saved metadata to {metadata_file}")

            # 5. Process clips in parallel using ThreadPoolExecutor
            from concurrent.futures import ThreadPoolExecutor

            swears_path = "/home/ren/viral/brainrotinator/assets/swears.txt"
            nextcloud_dir = os.environ.get(
                "NEXTCLOUD_TARGET",
                "/home/ren/nextcloud/html/data/renato97/files/viral_clips"
            )
            os.makedirs(nextcloud_dir, exist_ok=True)

            def normalize_for_mobile(path):
                temp = path + ".tmp.mp4"
                cmd = [
                    'ffmpeg', '-y', '-i', path,
                    '-c:v', 'libx264', '-preset', 'fast', '-crf', '23',
                    '-pix_fmt', 'yuv420p', '-profile:v', 'high', '-level:v', '4.1',
                    '-movflags', '+faststart',
                    '-c:a', 'aac', '-b:a', '128k',
                    temp
                ]
                subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
                os.replace(temp, path)

            def process_clip(i, clip):
                start = clip['start']
                end = clip['end']
                print(f"\n🎬 [Clip {i+1}] Starting processing: {start}s - {end}s")
                print(f"   [Clip {i+1}] Title: {clip.get('video_title_for_youtube_short', 'No Title')}")
                
                # Cut clip
                clip_filename = f"{video_title}_clip_{i+1}.mp4"
                clip_temp_path = os.path.join(output_dir, f"temp_{clip_filename}")
                clip_final_path = os.path.join(output_dir, clip_filename)
                
                # ffmpeg cut & speed up
                if args.speed != 1.0:
                    setpts = 1.0 / args.speed
                    cut_command = [
                        'ffmpeg', '-y', 
                        '-ss', str(start), 
                        '-to', str(end), 
                        '-i', input_video,
                        '-filter_complex', f'[0:v]setpts={setpts:.6f}*PTS[v];[0:a]atempo={args.speed:.6f}[a]',
                        '-map', '[v]', '-map', '[a]',
                        '-c:v', 'libx264', '-crf', '18', '-preset', 'fast',
                        '-c:a', 'aac',
                        clip_temp_path
                    ]
                else:
                    cut_command = [
                        'ffmpeg', '-y', 
                        '-ss', str(start), 
                        '-to', str(end), 
                        '-i', input_video,
                        '-c:v', 'libx264', '-crf', '18', '-preset', 'fast',
                        '-c:a', 'aac',
                        clip_temp_path
                    ]
                subprocess.run(cut_command, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
                
                # Generate SRT for this clip and get mute ranges for profanity
                srt_path = os.path.join(output_dir, f"{video_title}_clip_{i+1}.srt")
                print(f"   📝 [Clip {i+1}] Generating subtitles and detecting profanity...")
                success_srt, mute_ranges = generate_srt(
                    transcript, start, end, srt_path,
                    max_chars=20, max_duration=2.0,
                    swears_path=swears_path, return_mute_ranges=True,
                    speed_factor=args.speed
                )
                
                # Process vertical, passing the relative mute ranges for profanity muting
                success = process_video_to_vertical(clip_temp_path, clip_final_path, mute_ranges=mute_ranges)
                
                if success:
                    print(f"   ✅ [Clip {i+1}] Vertical crop ready")
                    
                    # Burn subtitles into the vertical video
                    if success_srt and os.path.exists(srt_path):
                        print(f"   🎬 [Clip {i+1}] Burning subtitles with Bangers font...")
                        subbed_path = os.path.join(output_dir, f"sub_{clip_filename}")
                        try:
                            burn_subtitles(
                                video_path=clip_final_path,
                                srt_path=srt_path,
                                output_path=subbed_path,
                                alignment=2,
                                fontsize=24,
                                font_name="Bangers",
                                font_color="#FFFFFF",
                                border_color="#000000",
                                border_width=3,
                                bg_color="#E91E63",
                                bg_opacity=1.0,
                                fonts_dir="/home/ren/viral/openshorts/fonts"
                            )
                            if os.path.exists(subbed_path):
                                os.replace(subbed_path, clip_final_path)
                                print(f"   ✅ [Clip {i+1}] Subtitles burned")
                        except Exception as e:
                            print(f"   ⚠️ [Clip {i+1}] Subtitle burn failed: {e}")
                    
                    # Normalize for mobile
                    print(f"   📱 [Clip {i+1}] Normalizing for mobile playback...")
                    try:
                        normalize_for_mobile(clip_final_path)
                    except Exception as e:
                        print(f"   ⚠️ [Clip {i+1}] Normalization failed: {e}")

                    # Copy to Nextcloud
                    import shutil
                    nextcloud_dest = os.path.join(nextcloud_dir, clip_filename)
                    print(f"   📤 [Clip {i+1}] Copying to Nextcloud: {nextcloud_dest}...")
                    try:
                        shutil.copy2(clip_final_path, nextcloud_dest)
                        # Fix permissions for Nextcloud
                        subprocess.run(['chmod', '664', nextcloud_dest])
                        # Scan nextcloud files for user renato97
                        print(f"   🔄 [Clip {i+1}] Triggering Nextcloud file scan...")
                        subprocess.run([
                            'docker', 'exec', '-u', 'www-data', 'nextcloud',
                            'php', 'occ', 'files:scan', 'renato97'
                        ], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
                        print(f"   ✅ [Clip {i+1}] Nextcloud sync complete!")
                    except Exception as e:
                        print(f"   ⚠️ [Clip {i+1}] Nextcloud upload failed: {e}")
                
                # Clean up temp cut
                if os.path.exists(clip_temp_path):
                    os.remove(clip_temp_path)

            print(f"🚀 Processing {len(clips_data['shorts'])} clips in parallel (max 4 workers)...")
            with ThreadPoolExecutor(max_workers=4) as executor:
                futures = [executor.submit(process_clip, i, clip) for i, clip in enumerate(clips_data['shorts'])]
                for future in futures:
                    try:
                        future.result()
                    except Exception as e:
                        print(f"❌ Clip processing failed with exception: {e}")
                        import traceback
                        traceback.print_exc()

    # Clean up original if requested
    if args.url and not args.keep_original and os.path.exists(input_video):
        os.remove(input_video)
        print(f"🗑️  Cleaned up downloaded video.")

    total_time = time.time() - script_start_time
    print(f"\n⏱️  Total execution time: {total_time:.2f}s")