import cv2
import os
import base64
from PIL import Image
import io
import numpy as np
import tempfile
import pathlib
import subprocess
import time
import shutil
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options

def create_multimodal_prompt_for_qwen(text_prompt, image_paths=None, video_path=None, audio_path=None, messages=None, sampling_rate=16000, text_before=False):
    """
    Create a multimodal prompt for Qwen VL models using OpenAI API format.
    
    Args:
        text_prompt: The text part of the prompt
        image_paths: List of paths to images to include (optional)
        video_path: Path to a video to include (optional)
        audio_path: Path to an audio file to include (optional)
        messages: Existing messages to append to (optional)
        sampling_rate: Sampling rate for audio (optional)
    
    Returns:
        A formatted prompt suitable for OpenAI API with Qwen VL models
    """
    # Create a list of content items for the user message
    content = []
    
    if text_before: # Add text content
        content.append({"type": "text", "text": text_prompt})

    # Add image content if provided
    if image_paths:
        for image_path in image_paths:
            # Read the image file and encode it as base64
            with open(image_path, "rb") as image_file:
                base64_image = base64.b64encode(image_file.read()).decode("utf-8")
            
            content.append({
                "type": "image_url",
                "image_url": {
                    "url": f"data:image/jpeg;base64,{base64_image}"
                }
            })
            print("Added img")
    
    # Add video content if provided
    if video_path and os.path.exists(video_path):
        try:
            # Read the video file and encode it as base64
            with open(video_path, "rb") as video_file:
                base64_video = base64.b64encode(video_file.read()).decode("utf-8")
            # Add the video
            content.append({
                "type": "video_url",
                "video_url": {
                    "url": f"data:video/mp4;base64,{base64_video}",
                }
            })
            print("Added video")
        except Exception as e:
            print(f"Error adding video: {e}")

    # Add audio content if provided
    if audio_path and os.path.exists(audio_path):
        try:
            # Read the audio file and encode it as base64
            with open(audio_path, "rb") as audio_file:
                base64_audio = base64.b64encode(audio_file.read()).decode("utf-8")
            # Add the audio
            content.append({
                "type": "audio_url",
                "audio_url": {
                    "url": f"data:audio/wav;base64,{base64_audio}"
                }
            })
            print("Added audio")
        except Exception as e:
            print(f"Error adding audio: {e}")

    if not text_before: # Add text content
        content.append({"type": "text", "text": text_prompt})

    # Create the messages array
    if messages is None:
        messages = [
            {"role": "system", "content": "You are a helpful assistant that can analyze images and videos."},
            {"role": "user", "content": content}
        ]
    else:
        messages += [{"role": "user", "content": content}]
    
    return messages


def create_multimodal_relative_prompt_for_qwen(text_prompt, video_path1=None, video_path2=None, audio_path1=None, audio_path2=None, text1=None, text2=None, messages=None, sampling_rate=16000, text_before=False):
    """
    Create a multimodal prompt for Qwen VL models using OpenAI API format.
    
    Args:
        text_prompt: The text part of the prompt
        image_paths: List of paths to images to include (optional)
        video_path: Path to a video to include (optional)
        audio_path: Path to an audio file to include (optional)
        messages: Existing messages to append to (optional)
        sampling_rate: Sampling rate for audio (optional)
    
    Returns:
        A formatted prompt suitable for OpenAI API with Qwen VL models
    """
    # Create a list of content items for the user message
    content = []
    
    if text_before: # Add text content
        content.append({"type": "text", "text": text_prompt})
    
    # Add text1+video1+audio1 contents
    if text1:
        content.append({"type": "text", "text": text1})
    if video_path1 and os.path.exists(video_path1):
        try:
            # Read the video file and encode it as base64
            with open(video_path1, "rb") as video_file:
                base64_video = base64.b64encode(video_file.read()).decode("utf-8")
            # Add the video
            content.append({
                "type": "video_url",
                "video_url": {
                    "url": f"data:video/mp4;base64,{base64_video}",
                }
            })
            print("Added video1")
        except Exception as e:
            print(f"Error adding video: {e}")
    if audio_path1 and os.path.exists(audio_path1):
        try:
            # Read the audio file and encode it as base64
            with open(audio_path1, "rb") as audio_file:
                base64_audio = base64.b64encode(audio_file.read()).decode("utf-8")
            # Add the audio
            content.append({
                "type": "audio_url",
                "audio_url": {
                    "url": f"data:audio/wav;base64,{base64_audio}"
                }
            })
            print("Added audio1")
        except Exception as e:
            print(f"Error adding audio: {e}")

    # Add text2+video2+audio2 contents
    if text2:
        content.append({"type": "text", "text": text2})
    if video_path2 and os.path.exists(video_path2):
        try:
            # Read the video file and encode it as base64
            with open(video_path2, "rb") as video_file:
                base64_video = base64.b64encode(video_file.read()).decode("utf-8")
            # Add the video
            content.append({
                "type": "video_url",
                "video_url": {
                    "url": f"data:video/mp4;base64,{base64_video}",
                }
            })
            print("Added video2")
        except Exception as e:
            print(f"Error adding video: {e}")
    if audio_path2 and os.path.exists(audio_path2):
        try:
            # Read the audio file and encode it as base64
            with open(audio_path2, "rb") as audio_file:
                base64_audio = base64.b64encode(audio_file.read()).decode("utf-8")
            # Add the audio
            content.append({
                "type": "audio_url",
                "audio_url": {
                    "url": f"data:audio/wav;base64,{base64_audio}"
                }
            })
            print("Added audio2")
        except Exception as e:
            print(f"Error adding audio: {e}")

    if not text_before: # Add text content
        content.append({"type": "text", "text": text_prompt})

    # Create the messages array
    if messages is None:
        messages = [
            {"role": "system", "content": "You are a helpful assistant that can analyze images and videos."},
            {"role": "user", "content": content}
        ]
    else:
        messages += [{"role": "user", "content": content}]
    
    return messages

def create_multimodal_relative_multiround_prompt_for_qwen(text_prompt, video_path1=None, video_path2=None, audio_path1=None, audio_path2=None, text1=None, answer1=None, text2=None, answer2=None, messages=None, sampling_rate=16000):
    """
    Create a multimodal prompt for Qwen VL models using OpenAI API format.
    # user: Content 1
    # assistant: Content 1 description
    # user: Content 2
    # assistant: Content 2 description
    # user: Content 1 vs content 2, which is better?
    # assistant: ?

    Args:
        text_prompt: The text part of the prompt
        image_paths: List of paths to images to include (optional)
        video_path: Path to a video to include (optional)
        audio_path: Path to an audio file to include (optional)
        messages: Existing messages to append to (optional)
        sampling_rate: Sampling rate for audio (optional)
    
    Returns:
        A formatted prompt suitable for OpenAI API with Qwen VL models
    """
    # Create a list of content items for the user message
    
    # Create the messages array
    if messages is None:
        messages = [
            {"role": "system", "content": "You are a helpful assistant that can analyze images and videos."}
        ]
    
    # Prompt1: describe video1 contents (user Question and assistant answer)
    content = []
    if text1:
        content.append({"type": "text", "text": text1})
    if video_path1 and os.path.exists(video_path1):
        try:
            # Read the video file and encode it as base64
            with open(video_path1, "rb") as video_file:
                base64_video = base64.b64encode(video_file.read()).decode("utf-8")
            # Add the video
            content.append({
                "type": "video_url",
                "video_url": {
                    "url": f"data:video/mp4;base64,{base64_video}",
                }
            })
            print("Added video1")
        except Exception as e:
            print(f"Error adding video: {e}")
    if audio_path1 and os.path.exists(audio_path1):
        try:
            # Read the audio file and encode it as base64
            with open(audio_path1, "rb") as audio_file:
                base64_audio = base64.b64encode(audio_file.read()).decode("utf-8")
            # Add the audio
            content.append({
                "type": "audio_url",
                "audio_url": {
                    "url": f"data:audio/wav;base64,{base64_audio}"
                }
            })
            print("Added audio1")
        except Exception as e:
            print(f"Error adding audio: {e}")
    messages += [{"role": "user", "content": content}]
    assert answer1 is not None
    content = [{"type": "text", "text": answer1}]
    messages += [{"role": "assistant", "content": content}]

    # Prompt2: describe video1 contents (user Question and assistant answer)
    content = []
    if text2:
        content.append({"type": "text", "text": text2})
    if video_path2 and os.path.exists(video_path2):
        try:
            # Read the video file and encode it as base64
            with open(video_path2, "rb") as video_file:
                base64_video = base64.b64encode(video_file.read()).decode("utf-8")
            # Add the video
            content.append({
                "type": "video_url",
                "video_url": {
                    "url": f"data:video/mp4;base64,{base64_video}",
                }
            })
            print("Added video2")
        except Exception as e:
            print(f"Error adding video: {e}")
    if audio_path2 and os.path.exists(audio_path2):
        try:
            # Read the audio file and encode it as base64
            with open(audio_path2, "rb") as audio_file:
                base64_audio = base64.b64encode(audio_file.read()).decode("utf-8")
            # Add the audio
            content.append({
                "type": "audio_url",
                "audio_url": {
                    "url": f"data:audio/wav;base64,{base64_audio}"
                }
            })
            print("Added audio2")
        except Exception as e:
            print(f"Error adding audio: {e}")
    messages += [{"role": "user", "content": content}]
    assert answer2 is not None
    content = [{"type": "text", "text": answer2}]
    messages += [{"role": "assistant", "content": content}]


    content = [{"type": "text", "text": text_prompt}]
    messages += [{"role": "user", "content": content}]
    
    return messages

def use_pulseaudio_sink(display_num, monitor_source):
    """
    Use an existing Pulseaudio virtual sink for a specific display number.
    
    Args:
        display_num: The display number for the sink (e.g., 1 for virtual_sink_1)
        monitor_source: The monitor source name (e.g., virtual_sink_1.monitor)
    
    Returns:
        The audio input device name for FFmpeg (e.g., virtual_sink_1.monitor or default)
    """
    sink_name = f"virtual_sink_{display_num}"
    print(f"Using module-null-sink '{sink_name}'")

    # Verify the monitor source exists
    audio_input_device = monitor_source
    try:
        result = subprocess.run(['pactl', 'list', 'sources', 'short'], capture_output=True, text=True, check=True)
        if audio_input_device in result.stdout:
            print(f"Pulseaudio monitor source '{audio_input_device}' found for FFmpeg.")
        else:
            print(f"Warning: Monitor source '{audio_input_device}' not found. Falling back to 'default'.")
            audio_input_device = "default"
    except subprocess.CalledProcessError as e:
        print(f"Error checking Pulseaudio sources: {e.stderr}. Falling back to 'default'.")
        audio_input_device = "default"
    except Exception as e:
        print(f"Unexpected error during source verification: {e}. Falling back to 'default'.")
        audio_input_device = "default"

    return audio_input_device

def start_xvfb(display_num=None):
    """
    Starts the X virtual framebuffer (Xvfb) on a specified display number.
    This creates a virtual screen where Chrome can render.

    Args:
        display_num: The display number to use for Xvfb
        
    Returns:
        A tuple containing the Xvfb process and the display string
    """

    # Use the provided display_num or find a free one
    if display_num is None:
        display_num = 1
        lock_file = f'/tmp/.X{display_num}-lock'
        while os.path.exists(lock_file):
            display_num += 1
            lock_file = f'/tmp/.X{display_num}-lock'
    else:
        # Ensure the display_num is an integer
        display_num = int(display_num)

    xvfb_display = f":{display_num}"

    # Xvfb command and arguments
    xvfb_cmd = [
        "Xvfb", xvfb_display, 
        "-screen", "0", "1280x720x16", # Screen 0, 1280x720 resolution, 24-bit color depth
        "-ac", # Disable access control (allow any host to connect)
        "+extension", "GLX", # Enable GLX extension (for GPU rendering features, though disabled for Chrome usually)
        "+render", # Enable Render extension
        "-noreset", # Don't reset server after last client disconnects (keeps it running)
    ]
    
    print(f"Starting Xvfb on display {xvfb_display}...")
    # Popen allows the Xvfb process to run in the background
    xvfb_process = subprocess.Popen(xvfb_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    
    # Wait a bit for Xvfb to fully initialize before Chrome tries to connect
    time.sleep(3) 
    
    return xvfb_process, xvfb_display

def cleanup_xvfb(xvfb_process):
    """
    Cleans up the Xvfb process if it was started by this script.
    
    Args:
        xvfb_process: The Xvfb process to clean up
    """
    if xvfb_process and xvfb_process.poll() is None:
        try:
            print("Terminating Xvfb process...")
            xvfb_process.terminate()
            xvfb_process.wait(timeout=5)  # Give it a few seconds to exit gracefully
        except Exception as e:
            print(f"Error terminating Xvfb process: {e}")
            if xvfb_process.poll() is None:  # If still running
                print("Xvfb process did not terminate, attempting to kill...")
                xvfb_process.kill()  # Force kill if terminate fails

def setup_headless_chrome_with_recording(display=":99", display_num=None, chromium_path=None):
    """
    Configures Chrome options for running in a headless environment with
    virtual display and suitable settings for media playbook and recording.
    
    Args:
        display: The display string to use for Chrome
        display_num: The display number to use for creating a unique user data directory
                    and configuring the specific PulseAudio virtual sink
        
    Returns:
        Chrome options object configured for recording
    """
    
    chrome_options = Options()
    # Recommended arguments for headless Chrome on a server
    chrome_options.add_argument('--no-sandbox') # Essential for Docker/CI environments
    chrome_options.add_argument('--disable-dev-shm-usage') # Overcomes shared memory limitations
    chrome_options.add_argument('--disable-gpu') # Necessary for headless environments
    chrome_options.add_argument('--window-size=1280,720') # Set a fixed resolution for consistent recording
    chrome_options.add_argument('--start-maximized') # Ensure full window size

    # Create a unique user data directory for Chrome based on display_num
    if display_num is not None:
        user_data_dir = f"/tmp/chrome_user_data_{display_num}"
        os.makedirs(user_data_dir, exist_ok=True)
        chrome_options.add_argument(f'--user-data-dir={user_data_dir}')
    
    # IMPORTANT: User needs to ensure this path is correct for their Chromium binary
    chrome_options.binary_location = os.path.join(chromium_path,"chrome-linux64/chrome")
    
    # Set the virtual display for Chrome to use Xvfb
    chrome_options.add_argument(f'--display={display}')
    
    # Audio/Video specific settings for Chrome
    chrome_options.add_argument('--autoplay-policy=no-user-gesture-required') # Allows media to play without user interaction
    chrome_options.add_argument('--disable-features=AutoplayIgnoreWebAudio')
    chrome_options.add_argument('--no-first-run') # Skips initial setup screens
    chrome_options.add_argument('--disable-infobars') # Hides "Chrome is being controlled by automated test software" bar
    chrome_options.add_argument('--disable-extensions') # Disables browser extensions
    chrome_options.add_argument('--disable-plugins-discovery') # Prevents plugin detection
    chrome_options.add_argument('--disable-background-networking')
    chrome_options.add_argument('--disable-sync')
    chrome_options.add_argument('--disable-web-security')
    chrome_options.add_argument('--disable-translate')
    chrome_options.add_argument('--allow-running-insecure-content') # Allows HTTP content on HTTPS pages (if applicable)
    
    # Configure PulseAudio sink for this specific Chrome instance
    if display_num is not None:
        # Force Chrome to use PulseAudio and specify the virtual sink
        chrome_options.add_argument('--alsa-output-device=pulse')
        # Note: The actual sink selection is done via environment variable PULSE_SINK
        # which will be set in the record_content function
    
    # Enable browser logging
    chrome_options.set_capability('goog:loggingPrefs', {'browser': 'ALL'})
    
    return chrome_options

def record_video_only(width=1280, height=720, fps=30, duration=10, 
                      output_path="video_recording.mp4", xvfb_display=":99"):
    """
    Records only video from the Xvfb display using FFmpeg.
    """
    print(f"Recording video only...")
    
    record_cmd = [
        "ffmpeg", "-y",
        "-use_wallclock_as_timestamps", "1",
        "-f", "x11grab",
        "-thread_queue_size", "2048",
        "-framerate", str(fps),
        "-video_size", f"{width}x{height}",
        "-probesize", "128M",
        "-i", f"{xvfb_display}.0+0,0",
        "-t", str(duration),
        "-c:v", "libx264",
        "-pix_fmt", "yuv420p",
        "-preset", "ultrafast",
        "-r", str(fps),
        "-crf", "28",
        "-an",  # Disable audio
        "-fps_mode", "cfr",
        "-max_muxing_queue_size", "4096",
        "-threads", "auto",
        output_path
    ]
    
    print(f"Recording video with FFmpeg command: {' '.join(record_cmd)}")
    
    # Create a local environment copy instead of modifying the global environment
    env = os.environ.copy()
    env['DISPLAY'] = xvfb_display

    try:
        subprocess.run(record_cmd, check=True, env=env)
        print(f"Video recording completed: {output_path}")
        return True
    except subprocess.CalledProcessError as e:
        print(f"FFmpeg video recording failed with exit code {e.returncode}.")
        print(f"FFmpeg stdout:\n{e.stdout.decode() if hasattr(e, 'stdout') and e.stdout else 'No stdout'}")
        print(f"FFmpeg stderr:\n{e.stderr.decode() if hasattr(e, 'stderr') and e.stderr else 'No stderr'}")
        return False
    except FileNotFoundError:
        print("Error: 'ffmpeg' command not found.")
        return False

def record_audio_only(duration=10, output_path="audio_recording.wav", audio_input_device="default"):
    """
    Records only audio from the specified Pulseaudio device using FFmpeg.
    """
    print(f"Recording audio from Pulseaudio device: '{audio_input_device}'")
    
    record_cmd = [
        "ffmpeg", "-y",
        "-f", "pulse",
        "-thread_queue_size", "2048",
        "-i", audio_input_device,
        "-t", str(duration),
        "-c:a", "pcm_s16le",
        "-ar", "44100",
        "-ac", "2",
        "-vn",  # Disable video
        "-threads", "auto",
        output_path
    ]
    
    print(f"Recording audio with FFmpeg command: {' '.join(record_cmd)}")
    
    # Use a local environment copy
    env = os.environ.copy()

    try:
        subprocess.run(record_cmd, check=True, env=env)
        print(f"Audio recording completed: {output_path}")
        return True
    except subprocess.CalledProcessError as e:
        print(f"FFmpeg audio recording failed with exit code {e.returncode}.")
        print(f"FFmpeg stdout:\n{e.stdout.decode() if hasattr(e, 'stdout') and e.stdout else 'No stdout'}")
        print(f"FFmpeg stderr:\n{e.stderr.decode() if hasattr(e, 'stderr') and e.stderr else 'No stderr'}")
        return False
    except FileNotFoundError:
        print("Error: 'ffmpeg' command not found.")
        return False

def record_content(html_path, output_video_path, duration, fps, enable_audio=False, sampling_rate=16000, display_num=99, max_console_logs=5, monitor_source=None, server_root='', chromium_path=None):
    """
    Open the HTML content and record a video of it, optionally with sound.
    Also captures browser console logs for debugging.
    Uses temporary files during recording to prevent corruption on crashes.
    
    Args:
        html_path: Path to the HTML content file
        output_video_path: Path to save the output video
        duration: Duration of the recording in seconds
        fps: Frames per second for the video
        enable_audio: Whether to record audio
        sampling_rate: Sampling rate for audio recording
        display_num: Display number for Xvfb
    
    Returns:
        Tuple containing:
        - Path to the recorded video
        - Path to the console logs file
        - Boolean indicating if console logs show no error
    """
    assert chromium_path is not None
    width, height = 1280, 720
    print(f"Recording content from {html_path}")
    
    # Store the original DISPLAY environment variable to restore it later
    old_display = os.environ.get('DISPLAY')
    
    # Create temporary file paths to prevent corruption on crashes
    temp_video_path = output_video_path.replace(".mp4", "TEMP.mp4")
    audio_path = os.path.splitext(output_video_path)[0] + ".wav" if enable_audio else None
    temp_audio_path = audio_path.replace(".wav", "TEMP.wav") if audio_path else None
    
    # Initialize console logs collection
    console_logs = []
    
    driver = None
    xvfb_process = None
    
    try:
        # Step 1: Start Xvfb virtual display with the specified display_num
        xvfb_process, xvfb_display = start_xvfb(display_num=display_num)
        
        # Create a local environment copy instead of modifying the global environment
        local_env = os.environ.copy()
        local_env['DISPLAY'] = xvfb_display
        print(f"Using display {xvfb_display} for this session")

        # Step 2: Set up Pulseaudio server with a unique sink name based on display_num
        audio_input_device = None
        if enable_audio:
            audio_input_device = use_pulseaudio_sink(display_num=display_num, monitor_source=monitor_source)
            print(f"FFmpeg configured to capture audio from Pulseaudio device: '{audio_input_device}'")
        
        # Step 3: Configure Chrome with a unique user data directory and PulseAudio sink
        chrome_options = setup_headless_chrome_with_recording(xvfb_display, display_num=display_num, chromium_path=chromium_path)
        driver_path = os.path.join(chromium_path,"chromedriver-linux64/chromedriver")
        service = Service(driver_path)

        # Configure environment for Chrome to use the specific virtual sink
        chrome_env = local_env.copy()
        if enable_audio and display_num is not None:
            sink_name = f"virtual_sink_{display_num}"
            chrome_env['PULSE_SINK'] = sink_name
            print(f"Chrome configured to output audio to PulseAudio sink: '{sink_name}'")

        # Step 4: First pass - Record video only to temporary file
        print("Initializing Chrome driver for video recording...")
        # Pass environment variables directly to the Chrome service without modifying global environment
        if enable_audio and display_num is not None:
            service.env = chrome_env
        driver = webdriver.Chrome(options=chrome_options, service=service)
        print("Chrome driver initialized.")

        if server_root == '':
            file_uri = pathlib.Path(html_path).absolute().as_uri()
        else:
            html_path_abs = pathlib.Path(html_path).absolute()
            relative_html_path = html_path_abs.relative_to(server_root)
            file_uri = f"http://localhost:8000/{relative_html_path.as_posix()}"
        print(f"Navigating to: {file_uri}")
        driver.get(file_uri)
        # Wait a moment for the page to load
        time.sleep(3)
        
        # Automatically click the start-button to enable audio
        try:
            print("Attempting to click start-button...")
            driver.execute_script("document.getElementById('start-button') && document.getElementById('start-button').click();")
            print("Start button clicked or not found")
        except Exception as e:
            print(f"Error clicking start button: {e}")

        print(f"Recording video to temporary file: {temp_video_path}")
        success_video = record_video_only(
            width=width,
            height=height,
            fps=fps,
            duration=duration,
            output_path=temp_video_path,
            xvfb_display=xvfb_display
        )

        if success_video:
            print("Video recording completed successfully!")
            print(f"Temporary video file: {os.path.abspath(temp_video_path)}")
        else:
            print("Video recording failed.")

        # Collect console logs before quitting the driver
        try:
            browser_logs = driver.get_log('browser')
            for log in browser_logs:
                log_message = f"{log['level']}: {log['message']}"
                console_logs.append(log_message)
        except Exception as e:
            print(f"Error collecting console logs: {e}")

        # Step 5: Close Chrome
        print("Quitting Chrome driver...")
        driver.quit()
        driver = None

        # Step 6: Second pass - Record audio only if enabled to temporary file
        success_audio = True  # Default to True if audio is not enabled
        if enable_audio and audio_input_device:
            print("Initializing Chrome driver for audio recording...")
            # Pass environment variables directly to the Chrome service (same as video recording)
            if enable_audio and display_num is not None:
                service.env = chrome_env
            driver = webdriver.Chrome(options=chrome_options, service=service)
            print("Chrome driver initialized.")

            if server_root == '':
                file_uri = pathlib.Path(html_path).absolute().as_uri()
            else:
                html_path_abs = pathlib.Path(html_path).absolute()
                relative_html_path = html_path_abs.relative_to(server_root)
                file_uri = f"http://localhost:8000/{relative_html_path.as_posix()}"
            print(f"Navigating to: {file_uri}")
            driver.get(file_uri)
            # Wait a moment for the page to load
            time.sleep(3)
            
            # Automatically click the start-button to enable audio
            try:
                print("Attempting to click start-button...")
                driver.execute_script("document.getElementById('start-button') && document.getElementById('start-button').click();")
                print("Start button clicked or not found")
            except Exception as e:
                print(f"Error clicking start button: {e}")

            print(f"Recording audio to temporary file: {temp_audio_path}")
            success_audio = record_audio_only(
                duration=duration,
                output_path=temp_audio_path,
                audio_input_device=audio_input_device
            )

            if success_audio:
                print("Audio recording completed successfully!")
                print(f"Temporary audio file: {os.path.abspath(temp_audio_path)}")
            else:
                print("Audio recording failed.")
                
            # Quit the driver after audio recording
            print("Quitting Chrome driver...")
            driver.quit()
            driver = None

        # Step 7: Move temporary files to final locations only if recording was successful
        if success_video:
            print(f"Moving video from temporary file to final location: {output_video_path}")
            shutil.move(temp_video_path, output_video_path)
            print(f"Video file moved successfully: {os.path.abspath(output_video_path)}")
        else:
            # Clean up temporary video file if recording failed
            if os.path.exists(temp_video_path):
                os.remove(temp_video_path)
                print("Removed temporary video file due to recording failure")

        if enable_audio and success_audio and temp_audio_path:
            print(f"Moving audio from temporary file to final location: {audio_path}")
            shutil.move(temp_audio_path, audio_path)
            print(f"Audio file moved successfully: {os.path.abspath(audio_path)}")
        elif enable_audio and temp_audio_path:
            # Clean up temporary audio file if recording failed
            if os.path.exists(temp_audio_path):
                os.remove(temp_audio_path)
                print("Removed temporary audio file due to recording failure")
            
    except Exception as e:
        print(f"An unexpected error occurred during the recording process: {e}")
        # Print a detailed traceback for debugging unexpected errors
        import traceback
        traceback.print_exc()
        
        # Clean up temporary files on exception
        if os.path.exists(temp_video_path):
            try:
                os.remove(temp_video_path)
                print("Cleaned up temporary video file after exception")
            except Exception as cleanup_e:
                print(f"Error cleaning up temporary video file: {cleanup_e}")
        
        if temp_audio_path and os.path.exists(temp_audio_path):
            try:
                os.remove(temp_audio_path)
                print("Cleaned up temporary audio file after exception")
            except Exception as cleanup_e:
                print(f"Error cleaning up temporary audio file: {cleanup_e}")
        
    finally:
        # Quit the driver if it's still running
        if driver:
            try:
                print("Quitting Chrome driver...")
                driver.quit()
            except Exception as e:
                print(f"Error quitting Chrome driver: {e}")
        
        # Restore the original DISPLAY environment variable
        if old_display:
            # Don't modify the global environment variable
            print(f"Original DISPLAY was: {old_display}")
        
        # Clean up Xvfb process
        if xvfb_process:
            cleanup_xvfb(xvfb_process)
                
    # Save console logs to a file next to the video (remove duplicates and limit to first 10 lines)
    unique_logs = []
    for log in console_logs:
        if log not in unique_logs:
            unique_logs.append(log)
    
    # Take only the first 10 lines
    limited_logs = unique_logs[:max_console_logs]
    console_logs_text = "\n".join(limited_logs)
    
    console_logs_path = os.path.splitext(output_video_path)[0] + "_console_logs.txt"
    with open(console_logs_path, "w", encoding="utf-8") as f:
        f.write(console_logs_text if len(limited_logs) > 0 else "No console logs captured")
    print(f"Console logs saved to {console_logs_path} (first 10 unique lines)")
    
    return output_video_path, console_logs_path, len(console_logs)==0 or "SEVERE" not in console_logs
