import os

import numpy as np
import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, WhisperFeatureExtractor
from moviepy.editor import VideoFileClip, AudioFileClip
import nltk
nltk.download('punkt', quiet=True)
from nltk.tokenize import sent_tokenize


def transcribe(video_file, transcribe_to_text=True, transcribe_to_srt=True, target_language='en'):
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

    model_id = "openai/whisper-large-v3"
    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)
    feature_extractor = WhisperFeatureExtractor.from_pretrained(model_id)

    video = VideoFileClip(video_file)
    audio = video.audio
    duration = audio.duration
    chunk_duration = 60
    n_chunks = int(np.ceil(duration / chunk_duration))

    full_transcription = ""
    for i in range(n_chunks):
        start_time = i * chunk_duration
        end_time = min((i + 1) * chunk_duration, duration)

        audio_chunk = audio.subclip(start_time, end_time)
        temp_file_path = f"temp_audio_chunk_{i}.wav"
        audio_chunk.write_audiofile(temp_file_path, codec='pcm_s16le')

        sound_array = AudioFileClip(temp_file_path).to_soundarray(fps=16000)
        if sound_array.ndim > 1:
            sound_array = np.mean(sound_array, axis=1)

        input_features = feature_extractor(sound_array, sampling_rate=16000, return_tensors="pt").input_features
        input_features = input_features.to(device=device, dtype=torch_dtype)

        with torch.no_grad():
            if target_language:
                model.config.forced_decoder_ids = processor.get_decoder_prompt_ids(language=target_language,
                                                                                   task="transcribe")
            generated_ids = model.generate(input_features, max_length=448)
            transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]

        full_transcription += transcription + " "
        os.remove(temp_file_path)
        print(f"Processed chunk {i + 1}/{n_chunks}")

    # Split the transcription into sentences
    sentences = sent_tokenize(full_transcription.strip())

    # Estimate time for each sentence based on its length relative to the total transcription
    total_chars = sum(len(s) for s in sentences)
    sentence_times = []
    current_time = 0
    for sentence in sentences:
        sentence_duration = (len(sentence) / total_chars) * duration
        sentence_times.append((current_time, current_time + sentence_duration))
        current_time += sentence_duration

    output = ""
    if transcribe_to_text:
        output += "Text Transcription:\n" + full_transcription + "\n\n"

    if transcribe_to_srt:
        output += "SRT Transcription:\n"
        for i, (sentence, (start, end)) in enumerate(zip(sentences, sentence_times), 1):
            output += f"{i}\n{format_time(start)} --> {format_time(end)}\n{sentence}\n\n"

    return output

def format_time(seconds):
    m, s = divmod(seconds, 60)
    h, m = divmod(m, 60)
    return f"{int(h):02d}:{int(m):02d}:{s:06.3f}".replace('.', ',')