import json
import io
import cv2
from tempfile import TemporaryDirectory
import numpy as np
from PIL import Image
import os
import base64
import re
import time

def modify_options(dict):
    new_str = []
    for key, value in dict.items():
        new_str.append(f"{key}. {value}")
    return new_str

def clean_and_convert_to_dict(json_str):
    clean_str = json_str.replace('```json', '').replace('```', '').replace('\n', '').strip()
    try:
        json_dict = json.loads(clean_str)
        return json_dict
    except json.JSONDecodeError as e:
        print("JSON decoding error:", e)
        error_position = e.pos
        if "Expecting ',' delimiter" in str(e):
            clean_str = clean_str[:error_position] + ',' + clean_str[error_position:]
        elif "Expecting ':' delimiter" in str(e):
            clean_str = clean_str[:error_position] + ':' + clean_str[error_position:]

        try:
            json_dict = json.loads(clean_str)
            return json_dict
        except json.JSONDecodeError:
            return None

def random_extract_and_encode_frames(video_path, num_frames=10):
    """
    Extracts frames from a video at equally spaced intervals and return them as a list of base64-encoded strings.
    
    :param video_path: Path to the video file.
    :param num_frames: Total number of frames to extract.
    :return: List of base64-encoded strings representing the extracted frames.
    """
    encoded_frames = []
    
    with TemporaryDirectory() as output_dir:
        # Open the video file
        cap = cv2.VideoCapture(video_path)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        
        # Calculate the frame indexes to extract
        frame_indexes = [int(total_frames * i / (num_frames - 1)) for i in range(num_frames)]
        print(f"Frame indexes to extract: {frame_indexes}")
        frame_indexes[-1] -= 1
        for index, frame_index in enumerate(frame_indexes):
            # Go to the frame index
            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
            success, frame = cap.read()
            
            if not success:
                print(f"Error: Could not read frame at index {frame_index}.")
                continue
            
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            image = Image.fromarray(frame_rgb)
            file_path = os.path.join(output_dir, f"{index}.png")
            image.save(file_path)
            with open(file_path, "rb") as image_file:
                encoded_frame = base64.b64encode(image_file.read()).decode('utf-8')
                encoded_frames.append(encoded_frame)
        
        # Release the video capture object
        cap.release()
    
    return encoded_frames

def extract_and_encode_frames(video_path, frame_indexes):
    """
    Extract frames from a video at the specified indexes and return them as a list of base64-encoded strings.

    :param video_path: Path to the video file.
    :param frame_indexes: List of frame indices to extract and encode.
    :return: List of base64-encoded strings representing the extracted frames.
    """
    encoded_frames = []

    with TemporaryDirectory() as output_dir:
        # Open the video file
        cap = cv2.VideoCapture(video_path)
        if len(frame_indexes) <= 8:
            frame_indexes = [0] + frame_indexes + [int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) - 1]
        if len(frame_indexes) <= 9:
            frame_indexes.append(int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) // 2)
        ranked_frame_indexes = np.array(frame_indexes)
        ranked_frame_indexes = np.sort(ranked_frame_indexes)
        print(ranked_frame_indexes)
        for index, frame_index in enumerate(ranked_frame_indexes):
            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
            success, frame = cap.read()
            
            if not success:
                print(f"Error: Could not read frame at index {frame_index}.")
                continue
            
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            image = Image.fromarray(frame_rgb)
            file_path = os.path.join(output_dir, f"{index}.png")
            image.save(file_path)
            with open(file_path, "rb") as image_file:
                encoded_frame = base64.b64encode(image_file.read()).decode('utf-8')
                encoded_frames.append(encoded_frame)

        cap.release()

    return encoded_frames[:10]

def extract_base_filename(video_path):
    """
    Extracts the base filename from a video path using regular expression.
    
    :param video_path: Full path to the video file.
    :return: Base filename extracted from the path.
    """
    match = re.search(r'/([^/]+)\.mp4$', video_path)
    if match:
        return match.group(1)
    else:
        raise ValueError("The video path does not contain a valid .mp4 file name.")

def encode_keyframes(video_path):
    """
    Encodes keyframe images associated with a video file into base64 strings.
    
    :param video_path: Full path to the video file.
    :return: List of base64-encoded strings representing the keyframes.
    """
    encoded_frames = []
    image_dir = os.path.join(os.path.dirname(video_path), "selectedframes")
    base_filename = extract_base_filename(video_path)
    
    for i in range(10): 
        image_path = os.path.join(image_dir, f"{base_filename}_{i}.jpeg")
        if not os.path.exists(image_path):
            break 
        try:
            with Image.open(image_path) as img:
                img = img.convert("RGB")
                buffer = io.BytesIO()
                img.save(buffer, format="JPEG")
                buffer.seek(0)
                encoded_frame = base64.b64encode(buffer.read()).decode('utf-8')
                encoded_frames.append(encoded_frame)
        except IOError:
            print(f"Error: Could not read image file {image_path}")
    
    return encoded_frames

def load_keyframes_for_gemini(video_path):
    """
    Loads keyframe images associated with a video file into PIL Image objects.
    
    :param video_path: Full path to the video file.
    :return: List of PIL.Image objects representing the keyframes.
    """
    images = []
    image_dir = os.path.join(os.path.dirname(video_path), "selectedframes")
    base_filename = extract_base_filename(video_path)
    
    for i in range(10): 
        image_path = os.path.join(image_dir, f"{base_filename}_{i}.jpeg")
        if not os.path.exists(image_path):
            break
        try:
            img = Image.open(image_path)
            img = img.convert("RGB")
            images.append(img)
        except IOError:
            print(f"Error: Could not read image file {image_path}")
    print(type(images[0]))
    return images

def random_extract_keyframes_for_gemini(video_path, num_frames=10):
    images = []
    
    keyframe_dir = os.path.join(os.path.dirname(video_path), "Gemini_keyframes")
    os.makedirs(keyframe_dir, exist_ok=True)

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print("Error: Could not read.")
        return images

    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    
    frame_indexes = [int(total_frames * i / (num_frames - 1)) for i in range(num_frames)]
    frame_indexes[-1] -= 1
    
    for frame_index in frame_indexes:
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
        success, frame = cap.read()
        
        if not success:
            print(f"Error: Could not read frame {i}.")
            continue
        
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        img = Image.fromarray(frame_rgb)
        img_path = os.path.join(keyframe_dir, f"frame_{frame_index}.png")
        
        img.save(img_path)
        loaded_img = Image.open(img_path)
        images.append(loaded_img)
    
    cap.release()
    return images

def extract_base_filename(path):
    """ Helper function to extract the base filename without extension. """
    return os.path.splitext(os.path.basename(path))[0]

def extract_key_frames_for_Qwen(video_path, num_frames=10, scale=0.3):
    """
    Extracts key frames from the video, scales them, and saves them in a specified directory.
    
    :param video_path: Full path to the video file.
    :param num_frames: Number of key frames to extract (default is 10).
    :param scale: Scale factor for resizing images (default is 1.0, which keeps original size).
    :return: None
    """
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print("Error: Could not open video.")
        return

    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    frame_indexes = [max(int(total_frames * i / (num_frames - 1)) - 1, 0) for i in range(num_frames)]
    print(f"Frame indexes to extract: {frame_indexes}")

    base_filename = extract_base_filename(video_path)
    image_dir = os.path.join(os.path.dirname(video_path), "selectedframes")
    os.makedirs(image_dir, exist_ok=True) 

    image_uris = []
    for i, frame_index in enumerate(frame_indexes):
        file_name = f"{base_filename}_{i}_{frame_index}_{scale}.png"
        if not os.path.exists(os.path.join(image_dir, file_name)):   
            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
            success, frame = cap.read()
            if not success:
                print(f"Error: Could not read frame at index {frame_index}.")
                continue

            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            img = Image.fromarray(frame_rgb)
            if scale != 1.0:
                new_size = (int(img.width * scale), int(img.height * scale))
                img = img.resize(new_size, Image.LANCZOS)

            img.save(os.path.join(image_dir, file_name))
        image_uris.append(f"file://{os.path.join(image_dir, file_name)}")
        
    cap.release()
    print(image_uris)
    return image_uris

def load_image_file_uris_for_Qwen(video_path, scale=1.0):
    """
    Creates a list of dictionaries, each containing the URI of an image file associated with a video, scaled by a given factor.
    
    :param video_path: Full path to the video file.
    :param scale: Scale factor for resizing images (default is 1.0, which keeps original size).
    :return: List of dictionaries with the image URIs.
    """
    image_uris = []
    image_dir = os.path.join(os.path.dirname(video_path), "selectedframes")
    base_filename = extract_base_filename(video_path)
    
    for i in range(10):
        
        scaled_image_file = f"{base_filename}_{i}_{scale}.jpeg"
        scaled_image_path = os.path.join(image_dir, scaled_image_file)
        scaled_image_uri = f"file://{scaled_image_path}"
        
        if os.path.exists(scaled_image_path):
            image_uris.append(scaled_image_uri)
        else:
            original_image_file = f"{base_filename}_{i}.jpeg"
            original_image_path = os.path.join(image_dir, original_image_file)
            if not os.path.exists(original_image_path):
                break
            with Image.open(original_image_path) as img:
                new_size = (int(img.width * scale), int(img.height * scale))
                img = img.resize(new_size, Image.LANCZOS)
                
                with open(scaled_image_path, 'wb') as f:
                    img.save(f, 'JPEG')
                    f.flush()
                    os.fsync(f.fileno()) 
                time.sleep(0.1)  

                image_uris.append(scaled_image_uri)
    print(image_uris)
    return image_uris