# video_generator.py
import pandas as pd
import numpy as np
import cv2
from pathlib import Path
from tqdm import tqdm
import traceback
from moviepy import VideoFileClip, AudioFileClip, CompositeAudioClip # Aggiunto

# --- Constants ---
NS_TO_S = 1e9 # Conversion factor from nanoseconds to seconds

# --- Drawing Constants ---
GAZE_COLOR = (0, 0, 255)  # Red in BGR
GAZE_RADIUS = 15
GAZE_THICKNESS = 2
PIP_SCALE = 0.25

# --- YOLO Drawing Constants ---
YOLO_BOX_COLOR = (0, 255, 255) # Cyan in BGR
YOLO_TEXT_COLOR = (0, 255, 255)
YOLO_THICKNESS = 2

# --- Pupil Plot Constants ---
PUPIL_PLOT_HISTORY = 200
PUPIL_PLOT_WIDTH = 350
PUPIL_PLOT_HEIGHT = 150
PUPIL_BG_COLOR = (80, 80, 80)
PUPIL_COLORS = {"Left": (80, 80, 255), "Right": (80, 255, 80), "Mean": (255, 255, 255)}
BLINK_TEXT_COLOR = (0, 0, 255)

# --- Fragmentation Plot Constants ---
FRAG_PLOT_HISTORY = 200
FRAG_PLOT_WIDTH = 350
FRAG_PLOT_HEIGHT = 150
FRAG_BG_COLOR = (80, 80, 80)
FRAG_LINE_COLOR = (255, 200, 100) # Light Blue

# --- Event Text Constants ---
EVENT_TEXT_COLOR = (255, 255, 255) # White
EVENT_BG_COLOR = (0, 0, 0) # Black

# --- On Surface Text Constants (NEW) ---
ON_SURFACE_TEXT_COLOR = (120, 255, 120) # Light Green

def _overlay_transparent(background, overlay, x, y):
    """
    Sovrappone un'immagine (overlay) con canale alpha su uno sfondo.
    """
    background_width = background.shape[1]
    background_height = background.shape[0]
    h, w = overlay.shape[0], overlay.shape[1]

    if x >= background_width or y >= background_height:
        return background

    if x + w > background_width:
        w = background_width - x
        overlay = overlay[:, :w]
    if y + h > background_height:
        h = background_height - y
        overlay = overlay[:h]

    if overlay.shape[2] < 4:
        return background

    alpha = overlay[:, :, 3] / 255.0
    alpha = np.expand_dims(alpha, axis=2)

    b, g, r = overlay[:, :, 0], overlay[:, :, 1], overlay[:, :, 2]
    bgr = np.dstack([b, g, r])

    background_subsection = background[y:y+h, x:x+w]
    composite = bgr * alpha + background_subsection * (1.0 - alpha)
    background[y:y+h, x:x+w] = composite
    return background

def _prepare_data(data_dir: Path, un_enriched_mode: bool, options: dict):
    """
    Loads and synchronizes all necessary DataFrames for video generation.
    MODIFIED: Merges enriched gaze data if needed for 'on surface' text.
    """
    try:
        world_timestamps = pd.read_csv(data_dir / 'world_timestamps.csv').sort_values('timestamp [ns]')
        gaze_df = pd.read_csv(data_dir / 'gaze.csv').sort_values('timestamp [ns]')
    except FileNotFoundError as e:
        raise FileNotFoundError(f"Essential file not found: {e}. Cannot generate video.")

    merged_data = pd.merge_asof(
        world_timestamps,
        gaze_df[['timestamp [ns]', 'gaze x [px]', 'gaze y [px]']],
        on='timestamp [ns]',
        direction='nearest',
        tolerance=pd.Timedelta('50ms').value
    )
    
    # --- NUOVA LOGICA: Aggiungi dati enriched se necessario ---
    if not un_enriched_mode and (options.get('overlay_on_surface_text') or options.get('crop_and_correct_perspective')):
        try:
            gaze_enr_df = pd.read_csv(data_dir / 'gaze_enriched.csv').sort_values('timestamp [ns]')
            
            cols_to_merge_enr = []
            if 'gaze detected on surface' in gaze_enr_df.columns:
                cols_to_merge_enr.append('gaze detected on surface')
            if 'aoi_name' in gaze_enr_df.columns:
                cols_to_merge_enr.append('aoi_name')

            if cols_to_merge_enr:
                 merged_data = pd.merge_asof(
                    merged_data,
                    gaze_enr_df[['timestamp [ns]'] + cols_to_merge_enr],
                    on='timestamp [ns]',
                    direction='backward' # backward fill per lo stato 'on surface'
                 )
        except FileNotFoundError:
             print("WARNING: 'gaze_enriched.csv' not found. 'On Surface' text overlay disabled.")
             options['overlay_on_surface_text'] = False
    # --- FINE NUOVA LOGICA ---

    # Calculate Fragmentation (Gaze Speed) for the video
    if options.get('overlay_fragmentation_plot'):
        merged_data['gaze_speed_px_per_s'] = np.sqrt(merged_data['gaze x [px]'].diff()**2 + merged_data['gaze y [px]'].diff()**2) / (merged_data['timestamp [ns]'].diff() / NS_TO_S)

    if options.get('overlay_pupil_plot'):
        try:
            pupil_df = pd.read_csv(data_dir / '3d_eye_states.csv').sort_values('timestamp [ns]')
            cols_to_merge = []
            if 'pupil diameter left [mm]' in pupil_df.columns:
                cols_to_merge.append('pupil diameter left [mm]')
            if 'pupil diameter right [mm]' in pupil_df.columns:
                cols_to_merge.append('pupil diameter right [mm]')
            
            if cols_to_merge:
                pupil_df['pupil_diameter_mean'] = pupil_df[cols_to_merge].mean(axis=1)
                cols_to_merge.append('pupil_diameter_mean')

                merged_data = pd.merge_asof(
                    merged_data,
                    pupil_df[['timestamp [ns]'] + cols_to_merge],
                    on='timestamp [ns]',
                    direction='backward'
                )
        except FileNotFoundError:
            print("WARNING: '3d_eye_states.csv' not found. Pupil plot disabled.")

    try:
        blinks_df = pd.read_csv(data_dir / 'blinks.csv')
        merged_data['is_blinking'] = False
        for _, row in blinks_df.iterrows():
            merged_data.loc[(merged_data['timestamp [ns]'] >= row['start timestamp [ns]']) & (merged_data['timestamp [ns]'] <= row['end timestamp [ns]']), 'is_blinking'] = True
    except FileNotFoundError:
        print("WARNING: 'blinks.csv' not found. Blink overlay disabled.")

    if options.get('crop_and_correct_perspective'):
        try:
            surface_df = pd.read_csv(data_dir / 'surface_positions.csv').sort_values('timestamp [ns]')
            corner_cols = ['tl x [px]', 'tl y [px]', 'tr x [px]', 'tr y [px]', 
                           'br x [px]', 'br y [px]', 'bl x [px]', 'bl y [px]']
            merged_data = pd.merge_asof(
                merged_data,
                surface_df[['timestamp [ns]'] + corner_cols],
                on='timestamp [ns]',
                direction='backward'
            )
        except FileNotFoundError:
            print("WARNING: Perspective option is active, but 'surface_positions.csv' not found. Option disabled.")
            options['crop_and_correct_perspective'] = False

    return merged_data

def _draw_generic_plot(frame: np.ndarray, data_points: list, min_val: float, max_val: float, width: int, height: int, position: tuple, title: str, color: tuple, bg_color: tuple):
    """Generic function to draw a single-line plot on the frame."""
    if not data_points or max_val == min_val:
        return frame

    x_pos, y_pos = position
    plot_area = frame[y_pos:y_pos+height, x_pos:x_pos+width]
    bg = np.full(plot_area.shape, bg_color, dtype=np.uint8)
    res = cv2.addWeighted(plot_area, 0.5, bg, 0.5, 0)
    frame[y_pos:y_pos+height, x_pos:x_pos+width] = res

    cv2.putText(frame, title, (x_pos + 5, y_pos + 15), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1)

    if len(data_points) < 2:
        return frame
    
    points = [
        (
            x_pos + int((i / (len(data_points) -1)) * width),
            y_pos + height - int(((val - min_val) / (max_val - min_val)) * (height - 25)) - 10
        )
        for i, val in enumerate(data_points) if pd.notna(val)
    ]

    if len(points) > 1:
        cv2.polylines(frame, [np.array(points, dtype=np.int32)], isClosed=False, color=color, thickness=2)

    return frame


def _draw_pupil_plot(frame: np.ndarray, plot_data_dict: dict, min_val: float, max_val: float, width: int, height: int, position: tuple):
    """Draws a multi-line graph for pupil data with a legend."""
    if not plot_data_dict or max_val == min_val:
        return frame
        
    x_pos, y_pos = position
    plot_area = frame[y_pos:y_pos+height, x_pos:x_pos+width]
    
    bg = np.full(plot_area.shape, PUPIL_BG_COLOR, dtype=np.uint8)
    res = cv2.addWeighted(plot_area, 0.5, bg, 0.5, 0)
    frame[y_pos:y_pos+height, x_pos:x_pos+width] = res

    # Draw the legend
    legend_y = y_pos + 15
    for name, color in PUPIL_COLORS.items():
        if name in plot_data_dict:
            cv2.putText(frame, name, (x_pos + 5, legend_y), cv2.FONT_HERSHEY_SIMPLEX, 0.4, color, 1)
            legend_y += 15

    # Draw each plot line
    for name, data_points in plot_data_dict.items():
        if len(data_points) < 2: continue
        
        points = [
            (
                x_pos + int((i / (PUPIL_PLOT_HISTORY - 1)) * width),
                y_pos + height - int(((val - min_val) / (max_val - min_val)) * (height - 50)) - 10 # 50px margin for legend
            )
            for i, val in enumerate(data_points) if pd.notna(val)
        ]
        
        if len(points) > 1:
            cv2.polylines(frame, [np.array(points, dtype=np.int32)], isClosed=False, color=PUPIL_COLORS[name], thickness=2)

    return frame


def create_custom_video(data_dir: Path, output_dir: Path, subj_name: str, options: dict, un_enriched_mode: bool, selected_events: list = None):
    """
    Main function for creating the video with selected overlays.
    MODIFIED: Can trim video to selected events.
    """
    video_out_path = output_dir / options.get('output_filename', f'video_output_{subj_name}.mp4')
    # Nuovo: percorso temporaneo per il video senza audio
    temp_video_path = output_dir / f'temp_no_audio_{subj_name}.mp4'
    
    print("Loading and synchronizing data...")
    try:
        sync_data = _prepare_data(data_dir, un_enriched_mode, options)
        # Assicurati che l'indice sia utile per ricerche veloci
        sync_data.set_index('timestamp [ns]', inplace=True, drop=False)
        sync_data.sort_index(inplace=True)
        
    except FileNotFoundError as e:
        print(f"CRITICAL ERROR: {e}. Cannot generate video.")
        return

    external_vid_path = data_dir / 'external.mp4'
    cap_ext = cv2.VideoCapture(str(external_vid_path))
    if not cap_ext.isOpened():
        print(f"ERROR: Cannot open external video: {external_vid_path}")
        return

    cap_int = None
    if options.get('include_internal_cam'):
        internal_vid_path = data_dir / 'internal.mp4'
        if internal_vid_path.exists():
            cap_int = cv2.VideoCapture(str(internal_vid_path))
            if not cap_int.isOpened():
                print("WARNING: Cannot open internal video, PiP disabled.")
                options['include_internal_cam'] = False
        else:
            print("WARNING: Internal video not found, PiP disabled.")
            options['include_internal_cam'] = False


    total_frames = int(cap_ext.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = cap_ext.get(cv2.CAP_PROP_FPS)
    original_w = int(cap_ext.get(cv2.CAP_PROP_FRAME_WIDTH))
    original_h = int(cap_ext.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    out_w, out_h = (1280, 720) if options.get('crop_and_correct_perspective') else (original_w, original_h)

    # Scrittura del video su un file temporaneo senza audio
    writer = cv2.VideoWriter(str(temp_video_path), cv2.VideoWriter_fourcc(*'mp4v'), fps, (out_w, out_h))
    print(f"The output video will be saved to: {video_out_path}")

    # Setup for Event Text & Trimming
    events_df = pd.DataFrame()
    try:
        events_df = pd.read_csv(data_dir / 'events.csv').sort_values('timestamp [ns]')
        events_df['name'] = events_df['name'].astype(str).str.replace(r'[\\/]', '_', regex=True)
    except FileNotFoundError:
        print("WARNING: 'events.csv' not found. Event text/trimming overlay disabled.")
        options['overlay_event_text'] = False
        options['trim_to_events'] = False

    current_event_name = ""
    event_idx = 0
    
    # Attempt to load YOLO detection data if the option is enabled.
    yolo_detections = pd.DataFrame()
    yolo_class_map = {}
    if options.get('overlay_yolo'):
        print("YOLO overlay enabled. Attempting to load detection data...")
        yolo_cache_path = output_dir / 'yolo_detections_cache.csv'
        id_map_path = output_dir / 'class_id_map.csv'
        try:
            yolo_detections = pd.read_csv(yolo_cache_path)
            id_map = pd.read_csv(id_map_path)
            yolo_class_map = pd.Series(id_map.class_name.values, index=id_map.track_id).to_dict()
            print("Successfully loaded YOLO detection data and class map.")
        except FileNotFoundError:
            print(f"WARNING: YOLO overlay is ON, but a required file was not found: {yolo_cache_path} or {id_map_path}")
            options['overlay_yolo'] = False
    
    # Setup for pupil plot
    pupil_plot_data = {"Left": [], "Right": [], "Mean": []}
    pupil_min, pupil_max = 0, 1
    pupil_cols = { "Left": "pupil diameter left [mm]", "Right": "pupil diameter right [mm]", "Mean": "pupil_diameter_mean" }
    if options.get('overlay_pupil_plot') and any(col in sync_data.columns for col in pupil_cols.values()):
        all_pupil_data = pd.concat([sync_data[col] for col in pupil_cols.values() if col in sync_data.columns]).dropna()
        if not all_pupil_data.empty: pupil_min, pupil_max = all_pupil_data.min(), all_pupil_data.max()

    # Setup for fragmentation plot
    frag_plot_data = []
    frag_min, frag_max = 0, 1
    if options.get('overlay_fragmentation_plot') and 'gaze_speed_px_per_s' in sync_data.columns:
        all_frag_data = sync_data['gaze_speed_px_per_s'].dropna()
        if not all_frag_data.empty: frag_min, frag_max = 0, all_frag_data.quantile(0.99)


    # Setup for Dynamic Heatmap
    gaze_history_for_heatmap = []
    heatmap_window_size_frames = 0
    if options.get('overlay_dynamic_heatmap'):
        # Il valore dello slider viene passato tramite il dizionario 'options'
        heatmap_window_seconds = options.get('heatmap_window_seconds', 2.0) 
        heatmap_window_size_frames = int(heatmap_window_seconds * fps)
    # --- FINE BLOCCO ---

    gaze_history_for_heatmap = []
    heatmap_window_size_frames = 0
    if options.get('overlay_dynamic_heatmap'):
        # Il valore dello slider viene passato tramite il dizionario 'options'
        heatmap_window_seconds = options.get('heatmap_window_seconds', 2.0) 
        heatmap_window_size_frames = int(heatmap_window_seconds * fps)

    # --- NUOVO: Inizializzazione per la scia dello sguardo ---
    gaze_path_history = []
    gaze_path_length = 10 # Memorizza gli ultimi 10 punti

    # --- NUOVA LOGICA: DEFINIZIONE DEI SEGMENTI DI FRAME DA PROCESSARE ---
    frame_segments = []
    if options.get('trim_to_events') and selected_events and not events_df.empty:
        print(f"Trimming video to {len(selected_events)} selected events.")
        
        # Filtra gli eventi per quelli selezionati
        events_to_process = events_df[events_df['name'].isin(selected_events)].copy()
        
        # Trova il timestamp di fine per ogni evento (è l'inizio del successivo)
        events_to_process['end_ts'] = events_to_process['timestamp [ns]'].shift(-1)
        # Per l'ultimo evento, la fine è l'ultimo timestamp disponibile nei dati
        events_to_process.fillna({'end_ts': sync_data['timestamp [ns]'].max()}, inplace=True)

        for _, event in events_to_process.iterrows():
            start_ts, end_ts = event['timestamp [ns]'], event['end_ts']
            
            # Trova gli indici dei frame corrispondenti nel dataframe sincronizzato
            segment_data = sync_data.loc[start_ts:end_ts]
            if not segment_data.empty:
                # 'frame' è la colonna originale da world_timestamps.csv che corrisponde all'indice del frame video
                start_frame = int(segment_data.iloc[0]['frame'])
                end_frame = int(segment_data.iloc[-1]['frame'])
                frame_segments.append((start_frame, end_frame))
        
        if not frame_segments:
            print("WARNING: Could not find any video frames for the selected events. Generating full video instead.")
            frame_segments = [(0, min(total_frames, len(sync_data)))]
    else:
        # Se il trim non è attivo, processa l'intero video
        frame_segments = [(0, min(total_frames, len(sync_data)))]
    # --- FINE NUOVA LOGICA ---

    try:
        # Loop principale attraverso i segmenti di frame (sarà uno solo se il trim non è attivo)
        with tqdm(total=sum(end - start for start, end in frame_segments), desc="Generating Video") as pbar:
            for start_frame, end_frame in frame_segments:
                # Posiziona la testina di lettura del video all'inizio del segmento
                cap_ext.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
                if cap_int:
                    # Tenta di sincronizzare anche il video interno (potrebbe non essere perfetto)
                    cap_int.set(cv2.CAP_PROP_POS_FRAMES, start_frame)

                for frame_idx in range(start_frame, end_frame):
                    ret_ext, frame = cap_ext.read()
                    if not ret_ext: break
                    
                    frame_data = sync_data.iloc[frame_idx]
                    current_ts = frame_data['timestamp [ns]']

                    # --- MODIFICA CHIAVE: Disegno Heatmap con Trasparenza ---
                    if options.get('overlay_dynamic_heatmap'):
                        if pd.notna(frame_data.get('gaze x [px]')):
                            gaze_history_for_heatmap.append((int(frame_data['gaze x [px]']), int(frame_data['gaze y [px]'])))
                        if len(gaze_history_for_heatmap) > heatmap_window_size_frames:
                            gaze_history_for_heatmap.pop(0)

                        if len(gaze_history_for_heatmap) > 1:
                            intensity_map = np.zeros((frame.shape[0], frame.shape[1]), dtype=np.uint8)
                            for point in gaze_history_for_heatmap:
                                cv2.circle(intensity_map, point, radius=25, color=50, thickness=-1)
                            
                            intensity_map = cv2.blur(intensity_map, (91, 91))
                            heatmap_color = cv2.applyColorMap(intensity_map, cv2.COLORMAP_JET)
                            heatmap_rgba = cv2.cvtColor(heatmap_color, cv2.COLOR_BGR2BGRA)
                            heatmap_rgba[:, :, 3] = intensity_map
                            frame = _overlay_transparent(frame, heatmap_rgba, 0, 0)
                    # --- FINE MODIFICA CHIAVE ---

                    # --- AGGIUNGI QUESTO BLOCCO PRIMA DI TUTTI GLI ALTRI OVERLAY ---
                    # Overlay Heatmap Dinamica
                    if options.get('overlay_dynamic_heatmap'):
                        # Aggiungi il punto di sguardo corrente alla cronologia
                        if pd.notna(frame_data.get('gaze x [px]')):
                            gaze_history_for_heatmap.append((int(frame_data['gaze x [px]']), int(frame_data['gaze y [px]'])))
                        
                        # Mantieni la cronologia entro la dimensione della finestra
                        if len(gaze_history_for_heatmap) > heatmap_window_size_frames:
                            gaze_history_for_heatmap.pop(0)

                        if len(gaze_history_for_heatmap) > 1:
                            heatmap_img = np.zeros((frame.shape[0], frame.shape[1]), dtype=np.uint8)
                            for point in gaze_history_for_heatmap:
                                cv2.circle(heatmap_img, point, radius=20, color=20, thickness=-1)
                            
                            heatmap_img = cv2.blur(heatmap_img, (81, 81))
                            heatmap_color = cv2.applyColorMap(heatmap_img, cv2.COLORMAP_JET)
                            frame = cv2.addWeighted(frame, 0.6, heatmap_color, 0.4, 0)
                    # --- FINE BLOCCO ---
                    M = None
                    if options.get('crop_and_correct_perspective') and pd.notna(frame_data.get('tl x [px]')):
                        src_pts = np.float32([[frame_data[c] for c in [f'{p} x [px]', f'{p} y [px]']] for p in ['tl','tr','br','bl']])
                        dst_pts = np.float32([[0, 0], [out_w, 0], [out_w, out_h], [0, out_h]])
                        M = cv2.getPerspectiveTransform(src_pts, dst_pts)
                        frame = cv2.warpPerspective(frame, M, (out_w, out_h))
                    elif frame.shape[1] != out_w or frame.shape[0] != out_h:
                        frame = cv2.resize(frame, (out_w, out_h))

                    # --- OVERLAYS ---

                    # Check for current event
                    if options.get('overlay_event_text') and event_idx < len(events_df):
                        # Trova l'ultimo evento iniziato prima del timestamp corrente
                        active_events = events_df[events_df['timestamp [ns]'] <= current_ts]
                        if not active_events.empty:
                            current_event_name = active_events.iloc[-1]['name']

                    if options.get('include_internal_cam') and cap_int is not None:
                        ret_int, frame_int = cap_int.read()
                        if ret_int:
                            pip_h = int(out_h * PIP_SCALE)
                            pip_w = int(frame_int.shape[1] * (pip_h / frame_int.shape[0]))
                            frame[10:10+pip_h, 10:10+pip_w] = cv2.resize(frame_int, (pip_w, pip_h))

                    # Draw YOLO object detection overlays
                    if options.get('overlay_yolo') and not yolo_detections.empty:
                        detections_for_frame = yolo_detections[yolo_detections['frame_idx'] == frame_idx]
                        for _, det in detections_for_frame.iterrows():
                            x1, y1, x2, y2 = int(det['x1']), int(det['y1']), int(det['x2']), int(det['y2'])
                            if M is not None:
                                pts = np.float32([[x1, y1], [x2, y2]]).reshape(-1, 1, 2)
                                transformed_pts = cv2.perspectiveTransform(pts, M)
                                if transformed_pts is not None:
                                    x1, y1 = int(transformed_pts[0][0][0]), int(transformed_pts[0][0][1])
                                    x2, y2 = int(transformed_pts[1][0][0]), int(transformed_pts[1][0][1])
                            
                            cv2.rectangle(frame, (x1, y1), (x2, y2), YOLO_BOX_COLOR, YOLO_THICKNESS)
                            track_id = int(det['track_id'])
                            class_name = yolo_class_map.get(track_id, f"ID:{track_id}")
                            cv2.putText(frame, class_name, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, YOLO_TEXT_COLOR, 2)

                    if options.get('overlay_gaze') and pd.notna(frame_data.get('gaze x [px]')):
                        gaze_x, gaze_y = frame_data['gaze x [px]'], frame_data['gaze y [px]']
                        px, py = (int(gaze_x), int(gaze_y))
                        if M is not None:
                            gaze_pt_transformed = cv2.perspectiveTransform(np.array([[[gaze_x, gaze_y]]], dtype=np.float32), M)
                            if gaze_pt_transformed is not None:
                                px, py = int(gaze_pt_transformed[0][0][0]), int(gaze_pt_transformed[0][0][1])
                        if 0 <= px < out_w and 0 <= py < out_h:
                            cv2.circle(frame, (px, py), GAZE_RADIUS, GAZE_COLOR, GAZE_THICKNESS, cv2.LINE_AA)

                            # --- NUOVO: Logica per la scia dello sguardo ---
                            gaze_path_history.append((px, py))
                            if len(gaze_path_history) > gaze_path_length:
                                gaze_path_history.pop(0)
                            
                            if options.get('overlay_gaze_path', True) and len(gaze_path_history) > 1:
                                for i in range(1, len(gaze_path_history)):
                                    thickness = int(np.ceil((i / len(gaze_path_history)) * (GAZE_THICKNESS)))
                                    cv2.line(frame, gaze_path_history[i-1], gaze_path_history[i], GAZE_COLOR, thickness, cv2.LINE_AA)
                            # --- FINE BLOCCO ---
                    
                    if frame_data.get('is_blinking', False):
                        cv2.putText(frame, "BLINK", (out_w - 150, out_h - 20), cv2.FONT_HERSHEY_TRIPLEX, 1.5, BLINK_TEXT_COLOR, 2)
                    
                    # --- NUOVO: Overlay testo 'On Surface' ---
                    if options.get('overlay_on_surface_text'):
                        # Priorità all'AOI
                        if pd.notna(frame_data.get('aoi_name')):
                            aoi_name = frame_data['aoi_name']
                            cv2.putText(frame, f"on AOI: {aoi_name}", (20, out_h - 20), cv2.FONT_HERSHEY_SIMPLEX, 0.7, ON_SURFACE_TEXT_COLOR, 2)
                        # Altrimenti, controlla la superficie generica
                        elif frame_data.get('gaze detected on surface') == True:
                            cv2.putText(frame, "on enrichment area", (20, out_h - 20), cv2.FONT_HERSHEY_SIMPLEX, 0.7, ON_SURFACE_TEXT_COLOR, 2)

                    if options.get('overlay_pupil_plot'):
                        for name, col in pupil_cols.items():
                            if col in frame_data and pd.notna(frame_data[col]): pupil_plot_data[name].append(frame_data[col])
                            if len(pupil_plot_data[name]) > PUPIL_PLOT_HISTORY: pupil_plot_data[name].pop(0)
                        frame = _draw_pupil_plot(frame, pupil_plot_data, pupil_min, pupil_max, PUPIL_PLOT_WIDTH, PUPIL_PLOT_HEIGHT, (out_w - PUPIL_PLOT_WIDTH - 10, 10))

                    if options.get('overlay_fragmentation_plot'):
                        if 'gaze_speed_px_per_s' in frame_data and pd.notna(frame_data['gaze_speed_px_per_s']): frag_plot_data.append(frame_data['gaze_speed_px_per_s'])
                        if len(frag_plot_data) > FRAG_PLOT_HISTORY: frag_plot_data.pop(0)
                        y_pos = (PUPIL_PLOT_HEIGHT + 20) if options.get('overlay_pupil_plot') else 10
                        frame = _draw_generic_plot(frame, frag_plot_data, frag_min, frag_max, FRAG_PLOT_WIDTH, FRAG_PLOT_HEIGHT, (out_w - FRAG_PLOT_WIDTH - 10, y_pos), "Fragmentation", FRAG_LINE_COLOR, FRAG_BG_COLOR)
                    
                    # Draw Event Text Overlay
                    if options.get('overlay_event_text') and current_event_name:
                        font_scale = 1.0; font_thickness = 2
                        (text_width, text_height), baseline = cv2.getTextSize(current_event_name, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness)
                        text_x = (out_w - text_width) // 2; text_y = 40
                        
                        rect_start = (text_x - 10, text_y - text_height - 5); rect_end = (text_x + text_width + 10, text_y + baseline)
                        sub_img = frame[rect_start[1]:rect_end[1], rect_start[0]:rect_end[0]]
                        black_rect = np.zeros(sub_img.shape, dtype=np.uint8)
                        res = cv2.addWeighted(sub_img, 0.5, black_rect, 0.5, 1.0)
                        frame[rect_start[1]:rect_end[1], rect_start[0]:rect_end[0]] = res
                        cv2.putText(frame, current_event_name, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, EVENT_TEXT_COLOR, font_thickness, cv2.LINE_AA)

                    writer.write(frame)
                    pbar.update(1) # Aggiorna la barra di avanzamento

    except Exception as e:
        print(f"An error occurred during video generation: {e}")
        traceback.print_exc()
    finally:
        print("Finalizing and releasing resources...")
        cap_ext.release()
        if cap_int:
            cap_int.release()
        writer.release()
        
        # --- NUOVA LOGICA: Aggiunta dell'audio ---
        if external_vid_path.exists():
            try:
                print("Adding audio to the generated video...")
                video_clip = VideoFileClip(str(temp_video_path))
                original_audio_clip = AudioFileClip(str(external_vid_path))
                
                # Taglia l'audio in base ai segmenti di trimming, se l'opzione è attiva
                if options.get('trim_to_events') and frame_segments:
                    start_frame, end_frame = frame_segments[0]
                    start_time = start_frame / fps
                    end_time = end_frame / fps
                    trimmed_audio = original_audio_clip.subclip(start_time, end_time)
                    final_clip = video_clip.set_audio(trimmed_audio)
                else:
                    final_clip = video_clip.set_audio(original_audio_clip)

                # Salva il video finale
                final_clip.write_videofile(str(video_out_path), codec='libx264', audio_codec='aac', logger=None)
                
                video_clip.close()
                original_audio_clip.close()
                if 'trimmed_audio' in locals():
                    trimmed_audio.close()
                
                # Rimuovi il file temporaneo senza audio
                temp_video_path.unlink(missing_ok=True)
                
                print("Video creation with audio completed!")
            except Exception as e:
                print(f"WARNING: An error occurred while adding audio. The video will be saved without audio. Error: {e}")
                # Rinominare il file temporaneo nel nome finale se l'aggiunta audio fallisce
                if temp_video_path.exists():
                    temp_video_path.rename(video_out_path)
                print("Video creation process completed, but without audio.")
        else:
            print("No external video file found to extract audio from. Video saved without audio.")
        # --- FINE NUOVA LOGICA ---