import gradio as gr
import os
import time
import sys
import io
import tempfile
import subprocess
import requests
from urllib.parse import urlparse
from pydub import AudioSegment
import logging
import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
import yt_dlp

class LogCapture(io.StringIO):
    def __init__(self, callback):
        super().__init__()
        self.callback = callback

    def write(self, s):
        super().write(s)
        self.callback(s)
        
logging.basicConfig(level=logging.INFO)

# Clone and install faster-whisper from GitHub
try:
    subprocess.run(["git", "clone", "https://github.com/SYSTRAN/faster-whisper.git"], check=True)
    subprocess.run(["pip", "install", "-e", "./faster-whisper"], check=True)
except subprocess.CalledProcessError as e:
    logging.error(f"Error during faster-whisper installation: {e}")
    sys.exit(1)

sys.path.append("./faster-whisper")

from faster_whisper import WhisperModel
from faster_whisper.transcribe import BatchedInferencePipeline

device = "cuda:0" if torch.cuda.is_available() else "cpu"

def download_audio(url, method_choice):
    parsed_url = urlparse(url)
    logging.info(f"Downloading audio from URL: {url} using method: {method_choice}")
    if parsed_url.netloc in ['www.youtube.com', 'youtu.be', 'youtube.com']:
        return download_youtube_audio(url, method_choice)
    else:
        return download_direct_audio(url, method_choice)

def download_youtube_audio(url, method_choice):
    methods = {
        'yt-dlp': youtube_dl_method,
        'pytube': pytube_method,
        'youtube-dl': youtube_dl_classic_method,
        'yt-dlp-alt': youtube_dl_alternative_method,
        'ffmpeg': ffmpeg_method,
        'aria2': aria2_method
    }
    method = methods.get(method_choice, youtube_dl_method)
    try:
        logging.info(f"Attempting to download YouTube audio using {method_choice}")
        return method(url)
    except Exception as e:
        logging.error(f"Error downloading using {method_choice}: {str(e)}")
        return None

def youtube_dl_method(url):
    logging.info("Using yt-dlp method")
    ydl_opts = {
        'format': 'bestaudio/best',
        'postprocessors': [{
            'key': 'FFmpegExtractAudio',
            'preferredcodec': 'mp3',
            'preferredquality': '192',
        }],
        'outtmpl': '%(id)s.%(ext)s',
    }
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        info = ydl.extract_info(url, download=True)
        logging.info(f"Downloaded YouTube audio: {info['id']}.mp3")
        return f"{info['id']}.mp3"

def pytube_method(url):
    logging.info("Using pytube method")
    from pytube import YouTube
    yt = YouTube(url)
    audio_stream = yt.streams.filter(only_audio=True).first()
    out_file = audio_stream.download()
    base, ext = os.path.splitext(out_file)
    new_file = base + '.mp3'
    os.rename(out_file, new_file)
    logging.info(f"Downloaded and converted audio to: {new_file}")
    return new_file

def youtube_dl_classic_method(url):
    logging.info("Using youtube-dl classic method")
    ydl_opts = {
        'format': 'bestaudio/best',
        'postprocessors': [{
            'key': 'FFmpegExtractAudio',
            'preferredcodec': 'mp3',
            'preferredquality': '192',
        }],
        'outtmpl': '%(id)s.%(ext)s',
    }
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        info = ydl.extract_info(url, download=True)
        logging.info(f"Downloaded YouTube audio: {info['id']}.mp3")
        return f"{info['id']}.mp3"

def youtube_dl_alternative_method(url):
    logging.info("Using yt-dlp alternative method")
    ydl_opts = {
        'format': 'bestaudio/best',
        'postprocessors': [{
            'key': 'FFmpegExtractAudio', 
            'preferredcodec': 'mp3',
            'preferredquality': '192',
        }],
        'outtmpl': '%(id)s.%(ext)s',
        'no_warnings': True,
        'quiet': True,
        'no_check_certificate': True,
        'prefer_insecure': True,
    }
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        info = ydl.extract_info(url, download=True) 
        logging.info(f"Downloaded YouTube audio: {info['id']}.mp3")
        return f"{info['id']}.mp3"

def ffmpeg_method(url):
    logging.info("Using ffmpeg method")
    output_file = tempfile.mktemp(suffix='.mp3')
    command = ['ffmpeg', '-i', url, '-vn', '-acodec', 'libmp3lame', '-q:a', '2', output_file]
    subprocess.run(command, check=True, capture_output=True)
    logging.info(f"Downloaded and converted audio to: {output_file}")
    return output_file

def aria2_method(url):  
    logging.info("Using aria2 method")
    output_file = tempfile.mktemp(suffix='.mp3')
    command = ['aria2c', '--split=4', '--max-connection-per-server=4', '--out', output_file, url]
    subprocess.run(command, check=True, capture_output=True)
    logging.info(f"Downloaded audio to: {output_file}")
    return output_file

def download_direct_audio(url, method_choice):
    logging.info(f"Downloading direct audio from: {url} using method: {method_choice}")
    if method_choice == 'wget':
        return wget_method(url)
    else:
        try:
            response = requests.get(url)
            if response.status_code == 200:
                with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as temp_file:
                    temp_file.write(response.content)
                    logging.info(f"Downloaded direct audio to: {temp_file.name}")
                    return temp_file.name
            else:
                raise Exception(f"Failed to download audio from {url}")
        except Exception as e:
            logging.error(f"Error downloading direct audio: {str(e)}")
            return None

def wget_method(url):
    logging.info("Using wget method")
    output_file = tempfile.mktemp(suffix='.mp3')  
    command = ['wget', '-O', output_file, url]
    subprocess.run(command, check=True, capture_output=True)
    logging.info(f"Downloaded audio to: {output_file}")
    return output_file

def trim_audio(audio_path, start_time, end_time):
    logging.info(f"Trimming audio from {start_time} to {end_time}")
    audio = AudioSegment.from_file(audio_path)
    audio_duration = len(audio) / 1000  # Duration in seconds

    # Default start and end times if None
    if start_time is None:
        start_time = 0
    if end_time is None or end_time > audio_duration:
        end_time = audio_duration

    # Validate times
    if start_time < 0 or end_time < 0:
        raise ValueError("Start time and end time must be non-negative.")
    if start_time >= end_time:
        raise gr.Error("End time must be greater than start time.")        
    if start_time > audio_duration:
        raise ValueError("Start time exceeds audio duration.")

    trimmed_audio = audio[start_time * 1000:end_time * 1000]
    trimmed_audio_path = tempfile.mktemp(suffix='.wav')
    trimmed_audio.export(trimmed_audio_path, format="wav")
    logging.info(f"Trimmed audio saved to: {trimmed_audio_path}")
    return trimmed_audio_path

def save_transcription(transcription):
    file_path = tempfile.mktemp(suffix='.txt')
    with open(file_path, 'w') as f:
        f.write(transcription)
    logging.info(f"Transcription saved to: {file_path}")
    return file_path

def get_model_options(pipeline_type):
    if pipeline_type == "faster-batched":
        return ["cstr/whisper-large-v3-turbo-int8_float32", "deepdml/faster-whisper-large-v3-turbo-ct2", "Systran/faster-whisper-large-v3", "GalaktischeGurke/primeline-whisper-large-v3-german-ct2"]
    elif pipeline_type == "faster-sequenced":
        return ["cstr/whisper-large-v3-turbo-int8_float32", "deepdml/faster-whisper-large-v3-turbo-ct2", "Systran/faster-whisper-large-v3", "GalaktischeGurke/primeline-whisper-large-v3-german-ct2"]
    elif pipeline_type == "transformers":
        return ["openai/whisper-large-v3", "openai/whisper-large-v3-turbo", "primeline/whisper-large-v3-german"]
    else:
        return []

def transcribe_audio(input_source, pipeline_type, model_id, dtype, batch_size, download_method, start_time=None, end_time=None, verbose=False):
    try:
        # Determine if input_source is a URL or file
        if isinstance(input_source, str):
            if input_source.startswith('http://') or input_source.startswith('https://'):
                audio_path = download_audio(input_source, download_method)
                # Handle potential errors during download
                if not audio_path or audio_path.startswith("Error"):
                    yield f"Error: {audio_path}", "", None
                    return
        else:
            # Assume input_source is an uploaded file object
            audio_path = input_source.name
            logging.info(f"Using uploaded audio file: {audio_path}")
            
    try:
        logging.info(f"Transcription parameters: pipeline_type={pipeline_type}, model_id={model_id}, dtype={dtype}, batch_size={batch_size}, download_method={download_method}")
        verbose_messages = f"Starting transcription with parameters:\nPipeline Type: {pipeline_type}\nModel ID: {model_id}\nData Type: {dtype}\nBatch Size: {batch_size}\nDownload Method: {download_method}\n"

        if verbose:
            yield verbose_messages, "", None

        if pipeline_type == "faster-batched":
            model = WhisperModel(model_id, device="auto", compute_type=dtype)
            pipeline = BatchedInferencePipeline(model=model)
        elif pipeline_type == "faster-sequenced":
            model = WhisperModel(model_id)
            pipeline = model.transcribe
        elif pipeline_type == "transformers":
            torch_dtype = torch.float16 if dtype == "float16" else torch.float32
            model = AutoModelForSpeechSeq2Seq.from_pretrained(
                model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
            )
            model.to(device)
            processor = AutoProcessor.from_pretrained(model_id)
            pipeline = pipeline(
                "automatic-speech-recognition",
                model=model,
                tokenizer=processor.tokenizer,
                feature_extractor=processor.feature_extractor,
                chunk_length_s=30,
                batch_size=batch_size,
                return_timestamps=True,
                torch_dtype=torch_dtype,
                device=device,
            )
        else:
            raise ValueError("Invalid pipeline type")

        if isinstance(input_source, str) and (input_source.startswith('http://') or input_source.startswith('https://')):
            audio_path = download_audio(input_source, download_method)
            verbose_messages += f"Audio file downloaded: {audio_path}\n"
            if verbose:
                yield verbose_messages, "", None

            if not audio_path or audio_path.startswith("Error"):
                yield f"Error: {audio_path}", "", None
                return
        else:
            audio_path = input_source

        start_time = float(start_time) if start_time else None
        end_time = float(end_time) if end_time else None

        if start_time is not None or end_time is not None:
            trimmed_audio_path = trim_audio(audio_path, start_time, end_time)
            audio_path = trimmed_audio_path
            verbose_messages += f"Audio trimmed from {start_time} to {end_time}\n"
            if verbose:
                yield verbose_messages, "", None

        start_time_perf = time.time()
        if pipeline_type in ["faster-batched", "faster-sequenced"]:
            segments, info = pipeline(audio_path, batch_size=batch_size)
        else:
            result = pipeline(audio_path)
            segments = result["chunks"]
        end_time_perf = time.time()

        transcription_time = end_time_perf - start_time_perf
        audio_file_size = os.path.getsize(audio_path) / (1024 * 1024)

        metrics_output = (
            f"Transcription time: {transcription_time:.2f} seconds\n"
            f"Audio file size: {audio_file_size:.2f} MB\n"
        )

        if verbose:
            yield verbose_messages + metrics_output, "", None

        transcription = ""

        for segment in segments:
            transcription_segment = (
                f"[{segment.start:.2f}s -> {segment.end:.2f}s] {segment.text}\n"
                if pipeline_type in ["faster-batched", "faster-sequenced"] else
                f"[{segment['timestamp'][0]:.2f}s -> {segment['timestamp'][1]:.2f}s] {segment['text']}\n"
            )
            transcription += transcription_segment
            if verbose:
                yield verbose_messages + metrics_output, transcription, None

        transcription_file = save_transcription(transcription)
        yield verbose_messages + metrics_output, transcription, transcription_file

    except Exception as e:
        logging.error(f"An error occurred during transcription: {str(e)}")
        yield f"An error occurred: {str(e)}", "", None

    finally:
        # Remove downloaded audio file
        if audio_path and os.path.exists(audio_path):
            os.remove(audio_path)
        # Remove trimmed audio file
        if 'trimmed_audio_path' in locals() and os.path.exists(trimmed_audio_path):
            os.remove(trimmed_audio_path)
        # Remove transcription file if needed
        if transcription_file and os.path.exists(transcription_file):
            os.remove(transcription_file)
            

with gr.Blocks() as iface:
    gr.Markdown("# Multi-Pipeline Transcription")
    gr.Markdown("Transcribe audio using multiple pipelines and models.")
    
    with gr.Row():
        input_source = gr.File(label="Audio Source (Upload a file or enter a URL/YouTube URL)")
        pipeline_type = gr.Dropdown(
            choices=["faster-batched", "faster-sequenced", "transformers"],
            label="Pipeline Type",
            value="faster-batched"
        )
        model_id = gr.Dropdown(
            label="Model",
            choices=get_model_options("faster-batched"),
            value=get_model_options("faster-batched")[0]
        )
    
    with gr.Row():
        dtype = gr.Dropdown(choices=["int8", "float16", "float32"], label="Data Type", value="int8")
        batch_size = gr.Slider(minimum=1, maximum=32, step=1, value=16, label="Batch Size")
        download_method = gr.Dropdown(
            choices=["yt-dlp", "pytube", "youtube-dl", "yt-dlp-alt", "ffmpeg", "aria2", "wget"],
            label="Download Method",
            value="yt-dlp"
        )
    
    with gr.Row():
        start_time = gr.Number(label="Start Time (seconds)", value=None, minimum=0)
        end_time = gr.Number(label="End Time (seconds)", value=None, minimum=0)
        verbose = gr.Checkbox(label="Verbose Output", value=True)  # Set to True by default
    
    transcribe_button = gr.Button("Transcribe")
    
    with gr.Row():
        metrics_output = gr.Textbox(label="Transcription Metrics and Verbose Messages", lines=10)
        transcription_output = gr.Textbox(label="Transcription", lines=10)
        transcription_file = gr.File(label="Download Transcription")
    
    def update_model_dropdown(pipeline_type):
        try:
            model_choices = get_model_options(pipeline_type)
            logging.info(f"Model choices for {pipeline_type}: {model_choices}")
        
            if model_choices:
                return gr.update(choices=model_choices, value=model_choices[0], visible=True)
            else:
                return gr.update(choices=["No models available"], value=None, visible=False)
        except Exception as e:
            logging.error(f"Error in update_model_dropdown: {str(e)}")
            return gr.update(choices=["Error"], value="Error", visible=True)
    
    #pipeline_type.change(update_model_dropdown, inputs=pipeline_type, outputs=model_id)
    pipeline_type.change(update_model_dropdown, inputs=[pipeline_type], outputs=model_id)
    
    def transcribe_with_progress(*args):
        for result in transcribe_audio(*args):
            yield result
    
    transcribe_button.click(
        transcribe_with_progress,
        inputs=[input_source, pipeline_type, model_id, dtype, batch_size, download_method, start_time, end_time, verbose],
        outputs=[metrics_output, transcription_output, transcription_file]
    )
    
    gr.Examples(
        examples=[
        ["https://www.youtube.com/watch?v=daQ_hqA6HDo", "faster-batched", "cstr/whisper-large-v3-turbo-int8_float32", "int8", 16, "yt-dlp", None, None, True],
        ["https://mcdn.podbean.com/mf/web/dir5wty678b6g4vg/HoP_453_-_The_Price_is_Right_-_Law_and_Economics_in_the_Second_Scholastic5yxzh.mp3", "faster-sequenced", "deepdml/faster-whisper-large-v3-turbo-ct2", "float16", 1, "ffmpeg", 0, 300, True],
        [None, "transformers", "openai/whisper-large-v3", "float16", 16, "yt-dlp", 60, 180, True]
        ],
        inputs=[input_source, pipeline_type, model_id, dtype, batch_size, download_method, start_time, end_time, verbose],
    )

iface.launch()