import os
import subprocess
from google.cloud import storage
import functions_framework
import logging

logger = logging.getLogger(__name__)


@functions_framework.cloud_event
def convert_video_to_audio(cloud_event):
    """Video to audio event-triggered cloud function."""

    data = cloud_event.data
    bucket_name = data['bucket']
    video_file_name = data['name']

    destination_bucket_name = os.environ.get("AUDIO_FILES_BUCKET_NAME")

    if not video_file_name.endswith(('.mp4', '.mov', '.avi', '.mkv')):
        logger.info(f"File {video_file_name} is not a supported video format.")
        return

    storage_client = storage.Client()
    bucket = storage_client.bucket(bucket_name)
    blob = bucket.blob(video_file_name)

    tmp_video_file = f"/tmp/{video_file_name}"
    blob.download_to_filename(tmp_video_file)

    audio_file_name = os.path.splitext(video_file_name)[0] + '.mp3'
    tmp_audio_file = f"/tmp/{audio_file_name}"

    command = f"ffmpeg -i {tmp_video_file} -vn -acodec libmp3lame -q:a 2 {tmp_audio_file}"
    subprocess.call(command, shell=True)

    destination_bucket = storage_client.bucket(destination_bucket_name)
    destination_blob = destination_bucket.blob(audio_file_name)
    destination_blob.upload_from_filename(tmp_audio_file)

    os.remove(tmp_video_file)
    os.remove(tmp_audio_file)

    logger.info(f"Converted {video_file_name} to {audio_file_name} and uploaded to {destination_bucket_name}.")

from google.cloud.speech_v2 import SpeechClient
from google.cloud.speech_v2.types import cloud_speech
from google.api_core.client_options import ClientOptions
from google.cloud import storage
import functions_framework
from typing import Dict
import logging
import time
import os

from . import chirp_model_long


logger = logging.getLogger(__file__)


def transcribe_batch_gcs(
    project_id: str,
    gcs_uri: str,
    region: str = "us-central1"
) -> cloud_speech.BatchRecognizeResults:
    """Transcribes audio from a Google Cloud Storage URI.

    Parameters
    ----------
        project_id: The Google Cloud project ID.
        gcs_uri: The Google Cloud Storage URI.

    Returns
    -------
        The RecognizeResponse.
    """

    client = SpeechClient(
        client_options=ClientOptions(
            api_endpoint=f"{region}-speech.googleapis.com",
        )
    )

    config = cloud_speech.RecognitionConfig(
        auto_decoding_config=cloud_speech.AutoDetectDecodingConfig(),
        language_codes=["fr-FR"],
        model="chirp",
    )

    file_metadata = cloud_speech.BatchRecognizeFileMetadata(uri=gcs_uri)

    request = cloud_speech.BatchRecognizeRequest(
        recognizer=f"projects/{project_id}/locations/{region}/recognizers/_",
        config=config,
        files=[file_metadata],
        recognition_output_config=cloud_speech.RecognitionOutputConfig(
            inline_response_config=cloud_speech.InlineOutputConfig(),
        ),
        processing_strategy=cloud_speech.BatchRecognizeRequest.ProcessingStrategy.DYNAMIC_BATCHING
    )

    operation = client.batch_recognize(request=request)

    logger.info("Waiting for operation to complete...")
    response = operation.result(timeout=1000)

    transcript = ""

    for result in response.results[gcs_uri].transcript.results:
        if len(result.alternatives) > 0:
            logger.info(f"Transcript: {result.alternatives[0].transcript}")
            transcript += f" \n{result.alternatives[0].transcript}"

    logger.debug(f"Transcript: {transcript}")

    return transcript


@functions_framework.cloud_event
def speech_to_text_transcription(cloud_event):
    """Audio file transcription via Speech-to-text API call."""

    data: Dict = cloud_event.data

    event_id = cloud_event["id"]
    event_type = cloud_event["type"]

    input_bucket = data["bucket"]
    audio_file_name = data["name"]

    destination_bucket_name = os.environ.get("RAW_TRANSCRIPTIONS_BUCKET_NAME")

    logger.info(f"Event ID: {event_id}")
    logger.info(f"Event type: {event_type}")
    logger.info(f"Bucket: {input_bucket}")
    logger.info(f"File: {audio_file_name}")

    storage_client = storage.Client()

    start = time.time()

    transcript = transcribe_batch_gcs(
        project_id=os.environ.get("PROJECT_ID"),
        gcs_uri=f"gs://{input_bucket}/{audio_file_name}"
    )

    stop = time.time()

    raw_transcription_file_name = os.path.splitext(audio_file_name)[0] + '_raw.txt'

    destination_bucket = storage_client.bucket(destination_bucket_name)
    destination_blob = destination_bucket.blob(raw_transcription_file_name)
    destination_blob.upload_from_string(transcript, content_type='text/plain; charset=utf-8')

    logger.debug(transcript)
    logger.info(f"JOB DONE IN {round(stop - start, 2)} SECONDS.")

import vertexai
from vertexai.generative_models import GenerativeModel

import functions_framework
from google.cloud import storage

import time
import logging
import os
from typing import Dict

logger = logging.getLogger(__file__)


@functions_framework.cloud_event
def transcription_correction(cloud_event):
    """Gemini API call to correct and enhance speech-to-text transcription."""

    data: Dict = cloud_event.data

    event_id = cloud_event["id"]
    event_type = cloud_event["type"]

    input_bucket = data["bucket"]
    raw_transcription_filename = data["name"]

    destination_bucket_name = os.environ.get("CURATED_TRANSCRIPTIONS_BUCKET_NAME")

    logger.info(f"Event ID: {event_id}")
    logger.info(f"Event type: {event_type}")
    logger.info(f"Bucket: {input_bucket}")
    logger.info(f"File: {raw_transcription_filename}")

    storage_client = storage.Client()

    input_bucket = storage_client.get_bucket(input_bucket)
    input_blob = input_bucket.get_blob(raw_transcription_filename)
    transcript = input_blob.download_as_string()

    vertexai.init(project=os.environ.get("PROJECT_ID"), location="us-central1")
    model = GenerativeModel(model_name="gemini-1.0-pro-002")

    prompt = f"""
        YOUR CUSTOM PROMPT GOES HERE. 
        PROVIDING CONTEXT AND GIVING INFORMATION ABOUT THE RESULT YOU EXPECT IS NECESSARY.

        {transcript}
    """

    n_tokens = model.count_tokens(prompt)

    logger.info(f"JOB : SPEECH-TO-TEXT TRANSCRIPTION CORRECTION. \n{n_tokens.total_billable_characters} BILLABLE CHARACTERS")
    logger.info(f"RESPONSE WILL PROCESS {n_tokens.total_tokens} TOKENS.")

    start = time.time()

    response = model.generate_content(prompt)

    stop = time.time()

    curated_filename = raw_transcription_filename.replace("_raw", "_curated")

    destination_bucket = storage_client.bucket(destination_bucket_name)
    destination_blob = destination_bucket.blob(curated_filename)
    destination_blob.upload_from_string(response.text, content_type='text/plain; charset=utf-8')

    logger.debug(response.text)
    logger.info(f"JOB DONE IN {round(stop - start, 2)} SECONDS.")

# locals.tf

locals {
    function_definitions = [
    {
        name = "convert_video_to_audio",
        source_dir = "../services/video_to_audio_cloud_function"
        input_bucket = var.video_files_bucket_name
    },
    {
        name = "speech_to_text_transcription",
        source_dir = "../services/transcript_cloud_function"
        input_bucket = var.audio_files_bucket_name
    },
    {
        name = "transcription_correction",
        source_dir = "../services/gemini_cloud_function"
        input_bucket = var.raw_transcriptions_bucket_name
    }
]
}

# cloud_functions.tf

resource "random_id" "bucket_prefix" {
  byte_length = 8
}

resource "google_storage_bucket" "source_code_bucket" {
  name                        = "${random_id.bucket_prefix.hex}-source-code-bucket"
  location                    = var.location
  force_destroy = true

  uniform_bucket_level_access = true
}

data "archive_file" "function_sources" {
  for_each    = { for def in local.function_definitions : def.name => def }
  type        = "zip"
  output_path = "/tmp/${each.value.name}-source.zip"
  source_dir  = each.value.source_dir
}

resource "google_storage_bucket_object" "function_sources" {
  for_each = data.archive_file.function_sources
  name     = "${basename(each.value.output_path)}-${each.value.output_md5}.zip"
  bucket   = google_storage_bucket.source_code_bucket.name
  source   = each.value.output_path
}

data "google_storage_project_service_account" "default" {}

resource "google_project_iam_member" "gcs_pubsub_publishing" {
  project = var.deploy_project
  role    = "roles/pubsub.publisher"
  member  = "serviceAccount:${data.google_storage_project_service_account.default.email_address}"
}

resource "google_service_account" "account" {
  account_id   = "gcf-sa"
  display_name = "Test Service Account - used for both the cloud function and eventarc trigger in the test"
}

resource "google_project_iam_member" "roles" {
  for_each = {
    "invoking"                   = "roles/run.invoker"
    "event_receiving"            = "roles/eventarc.eventReceiver"
    "artifactregistry_reader"    = "roles/artifactregistry.reader"
    "storage_object_admin"       = "roles/storage.objectUser"
    "speech_client"              = "roles/speech.client"
    "insights_collector_service" = "roles/storage.insightsCollectorService"
    "aiplatform_user"            = "roles/aiplatform.user"
  }
  project    = var.deploy_project
  role       = each.value
  member     = "serviceAccount:${google_service_account.account.email}"
  depends_on = [google_project_iam_member.gcs_pubsub_publishing]
}

resource "google_cloudfunctions2_function" "functions" {
  for_each = { for def in local.function_definitions : def.name => def }
  depends_on = [
    google_project_iam_member.roles["event_receiving"],
    google_project_iam_member.roles["artifactregistry_reader"],
  ]
  name        = each.value.name
  location    = var.location
  description = "Function to process ${each.value.name}"

  build_config {
    runtime     = "python39"
    entry_point = each.value.name
    environment_variables = {
      BUILD_CONFIG_TEST = "build_test"
    }
    source {
      storage_source {
        bucket = google_storage_bucket.source_code_bucket.name
        object = google_storage_bucket_object.function_sources[each.key].name
      }
    }
  }

  service_config {
    min_instance_count = 1
    max_instance_count = 3
    available_memory = "256M"
    timeout_seconds    = 60
    available_cpu = 4
    environment_variables = {
      PROJECT_ID                         = var.deploy_project
      AUDIO_FILES_BUCKET_NAME            = var.audio_files_bucket_name
      RAW_TRANSCRIPTIONS_BUCKET_NAME     = var.raw_transcriptions_bucket_name
      CURATED_TRANSCRIPTIONS_BUCKET_NAME = var.curated_transcriptions_bucket_name
    }
    ingress_settings               = "ALLOW_INTERNAL_ONLY"
    all_traffic_on_latest_revision = true
    service_account_email          = google_service_account.account.email
  }

  event_trigger {
    trigger_region        = var.location
    event_type            = "google.cloud.storage.object.v1.finalized"
    retry_policy          = "RETRY_POLICY_RETRY"
    service_account_email = google_service_account.account.email
    event_filters {
      attribute = "bucket"
      value     = google_storage_bucket.video_transcription_bucket_set[each.value.input_bucket].name
    }
  }
}

# buckets.tf

resource "google_storage_bucket" "video_transcription_bucket_set" {
  for_each = toset([
    var.video_files_bucket_name,
    var.audio_files_bucket_name,
    var.raw_transcriptions_bucket_name,
    var.curated_transcriptions_bucket_name
  ])
  name = each.value
  location = var.location
  storage_class = "STANDARD"
  force_destroy = true

  uniform_bucket_level_access = true
}