import runpod
from google.cloud import storage
import shutil
import os
import json
import uuid
from pathlib import Path
import logging
from datetime import timedelta 
import time

# Environment variable for GCS credentials (you should have this set)
gcs_credentials_json = os.getenv("GCS_SERVICE_ACCOUNT_JSON")

if gcs_credentials_json:
    credentials_dict = json.loads(gcs_credentials_json)  # Convert JSON string to Python dict
    client = storage.Client.from_service_account_info(credentials_dict)
else:
    raise ValueError("Missing GCS_SERVICE_ACCOUNT_JSON environment variable")

# Environment variable for the bucket name
bucket_name = os.getenv("GCS_BUCKET_NAME")

if not bucket_name:
    raise ValueError("Missing GCS_BUCKET_NAME environment variable")


# Function to upload a file to GCS
def upload_to_gcs(file_path, destination_blob_name):
    bucket = client.bucket(bucket_name)
    blob = bucket.blob(destination_blob_name)
    blob.upload_from_filename(file_path)
    # Return the public URL (optional)
    return blob.public_url

    
def generate_signed_url(blob_name, expiration_minutes=60):
    """Generate a signed URL for a GCS object."""
    bucket = client.bucket(bucket_name)
    blob = bucket.blob(blob_name)
    url = blob.generate_signed_url(expiration=timedelta(minutes=expiration_minutes))
    return url

def adjust_concurrency(current_concurrency):
    """
    Adjusts the concurrency level based on predefined conditions.
    """
    max_concurrency = 5
    min_concurrency = 1

    # Example condition: increase concurrency if it's below max
    if current_concurrency < max_concurrency:
        return current_concurrency + 1

    # Example condition: decrease concurrency if it's above min
    elif current_concurrency > min_concurrency:
        return current_concurrency - 1

    # Return the same concurrency if no changes are needed
    return current_concurrency

async def handler(job):
    try:
        video_blob_name = job['input'].get('video_path')
        audio_blob_name = job['input'].get('audio_path')

        # Ensure the paths are valid
        if not video_blob_name or not audio_blob_name:
            raise ValueError("Both video and audio paths must be provided with relative paths.")

        # uncomment once we done testing
        # Define temp paths (store in /workspace for persistence)
        workspace = Path("/tmp")
        video_local_path = workspace / "video.mp4"
        audio_local_path = workspace / "audio.m4a"
        output_filename = f"output_{uuid.uuid4().hex}.mp4"
        output_path = workspace / output_filename
        bucket = client.bucket(bucket_name)

        runpod.serverless.progress_update(job, f"Downloading Video")

        # Download video from GCS using the bucket name from the environment variable
        video_blob = bucket.blob(video_blob_name)
        video_blob.download_to_filename(video_local_path)

        runpod.serverless.progress_update(job, f"Downloading Audio")

        # Download audio from GCS using the same bucket
        audio_blob = bucket.blob(audio_blob_name)
        audio_blob.download_to_filename(audio_local_path)
        
        runpod.serverless.progress_update(job, f"Processing Video")

        # Dummy processing: Just copying the video file for now (you can replace this with actual processing logic)
        shutil.copy(str(video_local_path), str(output_path))

        runpod.serverless.progress_update(job, f"Uploading Video to GCS")

        # Upload the processed file back to GCS
        output_blob_name = f"light-test/{output_filename}"  # Save under light-test directory
        public_url = upload_to_gcs(str(output_path), output_blob_name)
        
        runpod.serverless.progress_update(job, f"Generating Signed URL")

        signed_url = generate_signed_url(output_blob_name)

        # Return the URL to the uploaded processed video
        return {
            "message": "Files processed and uploaded successfully",
            "signed_url": signed_url
        }

    except Exception as e:
        # Handle any exceptions and return a message
        return {
            "error": str(e)
        }


runpod.serverless.start(
    {
        "handler": handler, 
        "concurrency_modifier": adjust_concurrency,
    }
)