from google import genai
from google.genai import types
import csv
import os
import pandas as pd
import asyncio
import time

def upload_video(client, video_file_name):
    """Handles video upload and processing with status monitoring.
    
    Args:
        client: Cloud client instance with file management capabilities
        video_file_name (str): Local path to video file to upload
        
    Returns:
        File object: Processed video file reference with metadata
        
    Raises:
        ValueError: If video processing fails after upload
        
    Note:
        Implements polling loop to wait for cloud processing completion
    """
    # Upload video to cloud storage
    video_file = client.files.upload(file=video_file_name)
    # Wait for processing completion with periodic status checks
    while video_file.state == "PROCESSING":
        print('Waiting for video to be processed.')
        time.sleep(10)  # Check status every 10 seconds
        # Refresh file status from cloud
        video_file = client.files.get(name=video_file.name)
    # Handle processing failure
    if video_file.state == "FAILED":
        raise ValueError(video_file.state)
        
    # Return final processed file reference
    print(f'Video processing complete: ' + video_file.uri)
    return video_file

def append_to_csv(filename, data):
    """Appends a row to a CSV file.

    Args:
        filename (str): The name of the CSV file.
        data (list): A list of values to append as a new row.
    """
    try:
        with open(filename, 'a', newline='') as file:
            writer = csv.writer(file)
            writer.writerow(data)
    except Exception as e:
        print(f"An error occurred: {e}")

async def async_api_call(client, list_videos, model, filenname, prompt, idx):
    """Asynchronously processes a video through a generative AI model with specified parameters.
    
    Args:
        client: Cloud client instance for API communication
        list_videos (list): List of video inputs to process
        model (str): Name of the model to use for generation
        filenname (str): Reference filename for tracking
        prompt (str): Instruction prompt guiding the analysis
        idx (int): Index identifier for selecting video from list
        
    Returns:
        tuple: (filename, index, prompt, model response) containing:
            - filenname: Original filename reference
            - idx: Video index in input list
            - prompt: Analysis instruction used
            - response: Model's generated content response
    
    Note:
        Uses temperature=0 for deterministic outputs by disabling random sampling
    """
    response = client.models.generate_content(
        model=model,
        contents=[
            prompt,  # Analysis instructions for the model
            list_videos[idx],  # Selected video from input list
        ],
        config=types.GenerateContentConfig(
            temperature=0  # Zero temperature for consistent outputs
        )
    )
    return filenname, idx, prompt, response

async def main():

    # Configuration paths
    INTPHYS2_DATA_FOLDER = "" # Base data directory to replace
    YOUR_API_KEY = "" # Your Google GenAI API Key

    client = genai.Client(api_key=YOUR_API_KEY)

    # Define different analysis prompts
    prompt1 = "The video you’re seeing was generated by a simulator. Given how objects behave on earth, is the behaviour of the objects plausible? Your answer should be based on the events in the video and ignore the quality of the simulation. You should also ignore the framerate of the video, as it may be slowed down. Finish your answer by Yes or No"
    prompt2 = "Analyze the physics simulations in the video. Determine if the video depicts object behavior that is consistent with Earth's physical laws. Ignore simulation quality and playback speed. CRITICAL OUTPUT INSTRUCTION: Your response must be only a binary value indicating whether the object behavior is consistent with Earth's physical laws: 1: Object behavior is consistent with Earth's physical laws. 0: Object behavior is not consistent with Earth's physical laws. Please provide your analysis in the required format. Do not include any other text, explanation, context, preamble, or punctuation before or after the answer."
    prompt3 = "Analyze the physics simulations in the video. Determine if the video depicts object behavior that is consistent with Earth's physical laws. Ignore simulation quality and playback speed. CRITICAL OUTPUT INSTRUCTION: Your response must be only yes or no indicating whether the object behavior is consistent with Earth's physical laws: yes: Object behavior is consistent with Earth's physical laws. no: Object behavior is not consistent with Earth's physical laws. Please provide your analysis in the required format. Do not include any other text, explanation, context, preamble, or punctuation before or after the answer."
    # List of prompts for different analysis
    list_prompts = [prompt1, prompt2, prompt3]

    # Initialize model list
    list_models = ["gemini-1.5-pro", "gemini-2.5-flash-preview-04-17"]


    for split in ["Debug", "Test"]:
        # Load scene labels metadata
        file_labels = INTPHYS2_DATA_FOLDER+split+"/metadata.csv"
        df2 = pd.read_csv(file_labels)
        list_scenes_ids = df2["SceneIndex"].unique()
        # Go throug all models in the list
        for model in list_models:
            # Prepare result files for different prompts and seeds
            for idx_prompt, prompt in enumerate(list_prompts):
                # It's API call, so we cannot really control the seed but we observed different results between different runs even with the temperature equal to 0
                for seed in [1,2]:
                    label_filename = "Results/seed_"+str(seed)+"prompt"+str(idx_prompt)+"_"+split+"_"+model+".csv"
                    if not os.path.exists(label_filename):
                        with open(label_filename, "w+", newline="") as f:
                            # w = csv.DictWriter(f, ["filename", "response"])
                            w = csv.DictWriter(f, ["scene idx", "filename", "response", "ground_truth"])
                            w.writeheader()

        # Process each scene in the dataset
        for scene_idx in list_scenes_ids:
            data_to_read = df2[df2["SceneIndex"] == scene_idx]

            # Construct video file paths for different test cases
            filename_possible_video_1 = data_to_read[data_to_read["type"].isin(["1_Possible"])]['filename'].values[0]
            possible_video_1 = upload_video(INTPHYS2_DATA_FOLDER+split+'/Videos/'+filename_possible_video_1+'.mp4', seconds_per_frame=1)[:50]

            filename_impossible_video_1 = data_to_read[data_to_read["type"].isin(["1_Impossible"])]['filename'].values[0]
            impossible_video_1 = upload_video(INTPHYS2_DATA_FOLDER+split+'/Videos/'+filename_impossible_video_1+'.mp4', seconds_per_frame=1)[:50]

            filename_possible_video_2 = data_to_read[data_to_read["type"].isin(["2_Possible"])]['filename'].values[0]
            possible_video_2 = upload_video(INTPHYS2_DATA_FOLDER+split+'/Videos/'+filename_possible_video_2+'.mp4', seconds_per_frame=1)[:50]

            filename_impossible_video_2 = data_to_read[data_to_read["type"].isin(["2_Impossible"])]['filename'].values[0]
            impossible_video_2 = upload_video(INTPHYS2_DATA_FOLDER+split+'/Videos/'+filename_impossible_video_2+'.mp4', seconds_per_frame=1)[:50]

            # List of test cases with corresponding ground truths
            list_filenames = [filename_possible_video_1, filename_impossible_video_1, filename_possible_video_2, filename_impossible_video_2]
            list_videos = [possible_video_1, impossible_video_1, possible_video_2, impossible_video_2]
            ground_truth = [1,0,1,0]
            tasks = []
            # Process each video with different prompts and seeds
            for model in list_models:
                for idx in range(4): # Iterate through 4 test cases
                    for idx_prompt, prompt in enumerate(list_prompts):
                        for seed in [1,2]: # Two different random seeds
                            label_filename = "Results/seed_"+str(seed)+"prompt"+str(idx_prompt)+"_"+split+"_"+model+".csv"
                            # Make api call to get model's answer
                            task = async_api_call(client, list_videos, model, label_filename, prompt, idx)  # doesn't block until gather below
                            tasks.append(task)

            # Save collected responses to CSV    
            responses = await asyncio.gather(*tasks)
            for label_filename, idx, prompt, response in responses:
                append_to_csv(label_filename, [scene_idx, list_filenames[idx], response.text, ground_truth[idx]])

asyncio.run(main())
