import sys
import httpx
import os
import argparse
from datetime import datetime
from dotenv import load_dotenv
from deepgram import DeepgramClient, PrerecordedOptions
from deepgram_captions import DeepgramConverter, srt
from moviepy.video.io.VideoFileClip import VideoFileClip
from moviepy.audio.io.AudioFileClip import AudioFileClip
import deepl
import re

load_dotenv()

def cleanup_srt_punctuation(srt_content):
    # Split the SRT content into blocks
    blocks = re.split(r'\n\s*\n', srt_content.strip())
    
    parsed_blocks = []
    for block in blocks:
        lines = block.split('\n')
        if len(lines) >= 2:
            index = lines[0]
            timecode = lines[1]
            text = "\n".join(lines[2:]) if len(lines) > 2 else ""
            parsed_blocks.append({
                "index": index,
                "timecode": timecode,
                "text": text
            })

    # Rule 1: Clean up spaces before punctuation within each block
    for block in parsed_blocks:
        if "text" in block:
            block["text"] = re.sub(r'\s+([.,!?~:;。、])', r'\1', block["text"])

    # Rule 2 & 3: Handle leading punctuation and punctuation-only blocks
    for i in range(len(parsed_blocks)):
        block = parsed_blocks[i]
        if "text" not in block:
            continue
        
        text = block["text"].strip()
        
        # Check if the block is only punctuation
        if text and all(c in ".,!?~:;。、" or c.isspace() for c in text):
            for j in range(i - 1, -1, -1):
                prev_block = parsed_blocks[j]
                if "text" in prev_block and prev_block["text"].strip():
                    prev_block["text"] = prev_block["text"].rstrip() + " " + text
                    prev_block["text"] = re.sub(r'\s+([.,!?~:;。、])', r'\1', prev_block["text"])
                    break
            block["text"] = ""
            continue
            
        # Check if the block starts with leading punctuation (e.g. ", text")
        match = re.match(r'^([.,!?~:;。、\s]+)(.*)', block["text"])
        if match:
            lead_punct = match.group(1).strip()
            remaining_text = match.group(2)
            if lead_punct:
                for j in range(i - 1, -1, -1):
                    prev_block = parsed_blocks[j]
                    if "text" in prev_block and prev_block["text"].strip():
                        prev_block["text"] = prev_block["text"].rstrip() + " " + lead_punct
                        prev_block["text"] = re.sub(r'\s+([.,!?~:;。、])', r'\1', prev_block["text"])
                        break
                block["text"] = remaining_text

    # Reconstruct and re-index the SRT string, filtering out empty blocks
    reconstructed = []
    entry = 1
    for block in parsed_blocks:
        text = block["text"].strip()
        if text:
            reconstructed.append(f"{entry}\n{block['timecode']}\n{text}")
            entry += 1
            
    return "\n\n".join(reconstructed) + "\n"

def translate_srt_content(srt_content, deepl_api_key, target_lang):
    import deepl
    
    # Split the SRT content into blocks
    blocks = re.split(r'\n\s*\n', srt_content.strip())
    
    parsed_blocks = []
    text_list = []
    
    for block in blocks:
        lines = block.split('\n')
        if len(lines) >= 2:
            index = lines[0]
            timecode = lines[1]
            text = "\n".join(lines[2:]) if len(lines) > 2 else ""
            
            # Extract speaker tag if any (e.g. "[speaker 0] Hello" or "[Speaker 1]")
            tag = ""
            clean_text = text
            match = re.match(r'^(\[speaker \d+\]\s*)(.*)', text, re.IGNORECASE)
            if match:
                tag = match.group(1)
                clean_text = match.group(2)
                
            parsed_blocks.append({
                "index": index,
                "timecode": timecode,
                "tag": tag,
                "clean_text": clean_text
            })
            if clean_text.strip():
                text_list.append(clean_text)
        else:
            parsed_blocks.append({
                "raw": block
            })
            
    # Translate clean texts using DeepL text translation
    translator = deepl.Translator(deepl_api_key)
    translated_texts = []
    
    # Chunk text requests to avoid hitting DeepL payload size limits
    chunk_size = 50
    for i in range(0, len(text_list), chunk_size):
        chunk = text_list[i:i + chunk_size]
        try:
            results = translator.translate_text(chunk, target_lang=target_lang)
            translated_texts.extend([r.text for r in results])
        except Exception as e:
            print(f"Error translating chunk: {e}")
            translated_texts.extend(chunk)
            
    # Reassemble the parsed blocks
    text_idx = 0
    reconstructed = []
    entry = 1
    
    for block in parsed_blocks:
        if "raw" in block:
            reconstructed.append(block["raw"])
        else:
            clean_text = block["clean_text"]
            tag = block["tag"]
            
            if clean_text.strip():
                translated_text = translated_texts[text_idx] if text_idx < len(translated_texts) else clean_text
                text_idx += 1
                full_text = tag + translated_text
            else:
                full_text = tag + clean_text
                
            # Filter out empty blocks after translation and re-index sequentially
            stripped_text = full_text.strip()
            if stripped_text:
                reconstructed.append(f"{entry}\n{block['timecode']}\n{stripped_text}")
                entry += 1
                
    return "\n\n".join(reconstructed) + "\n"

def main():
    parser = argparse.ArgumentParser(description="Transcribe video/audio to SRT subtitles using Deepgram.")
    parser.add_argument("filepath", type=str, help="Path to the audio or video file to transcribe.")
    parser.add_argument("-m", "--model", type=str, default="nova-3", help="Deepgram model to use (default: %(default)s).")
    parser.add_argument("-l", "--language", type=str, default=None, help="BCP-47 language tag (e.g. 'en', 'es', 'fr'), or 'auto'/'detect' to enable automatic language detection.")
    parser.add_argument("--no-diarize", dest="diarize", action="store_false", help="Disable speaker diarization.")
    parser.add_argument("-t", "--translate-to", type=str, default=None, help="Translate the generated subtitles to this BCP-47 language tag (e.g. 'ko', 'en', 'ja') using DeepL.")
    parser.set_defaults(diarize=True)
    
    args = parser.parse_args()
    filepath = args.filepath

    # Resolve filepath. If it doesn't exist directly but exists in 'media/', use it from there.
    if not os.path.exists(filepath):
        media_fallback = os.path.join("media", filepath)
        if os.path.exists(media_fallback):
            filepath = media_fallback

    if not os.path.exists(filepath):
        print(f"Error: File '{filepath}' not found.")
        print("Please check the path or place the file in the 'media' directory.")
        return

    _, ext = os.path.splitext(filepath.lower())

    if ext == '.srt':
        if not args.translate_to:
            print("Error: When passing an .srt file, you must specify a target language using -t or --translate-to.")
            return

        deepl_api_key = os.getenv("DEEPL_API_KEY") or os.getenv("DEEPL_AUTH_KEY")
        if not deepl_api_key:
            print("Error: DEEPL_API_KEY or DEEPL_AUTH_KEY environment variable is not set.")
            print("Please set it in your environment or add it to your .env file to use translation.")
            return

        try:
            target_lang = args.translate_to.upper()
            if target_lang == "EN":
                target_lang = "EN-US"
            elif target_lang == "PT":
                target_lang = "PT-BR"

            base, _ = os.path.splitext(filepath)
            translated_srt_path = f"{base}.{args.translate_to.lower()}.srt"

            print(f"Translating {filepath} to {args.translate_to} using DeepL...")
            with open(filepath, "r", encoding="utf-8") as f:
                original_content = f.read()
                
            translated_content = translate_srt_content(original_content, deepl_api_key, target_lang)
            cleaned_content = cleanup_srt_punctuation(translated_content)
            
            with open(translated_srt_path, "w", encoding="utf-8") as f:
                f.write(cleaned_content)
                
            print(f"Successfully translated subtitles. Saved to: {translated_srt_path}")
        except Exception as translate_err:
            print(f"An error occurred during translation: {translate_err}")
        return

    api_key = os.getenv("DEEPGRAM_API_KEY")
    if not api_key:
        print("Error: DEEPGRAM_API_KEY environment variable is not set.")
        print("Please set it in your environment or add it to a .env file in the project directory.")
        return

    try:
        deepgram = DeepgramClient(api_key)
        is_audio = ext in {'.mp3', '.wav', '.m4a', '.flac', '.ogg', '.aac', '.wma', '.opus', '.webm', '.m4b', '.mp4a', '.aiff', '.aif', '.mp2'}

        audio_filepath = filepath
        should_remove_audio = False

        if not is_audio:
            audio_filepath = f"{filepath}-audio.mp3"
            should_remove_audio = False
            
            audio_exists = False
            if os.path.exists(audio_filepath) and os.path.getsize(audio_filepath) > 0:
                try:
                    with VideoFileClip(filepath) as video_clip:
                        video_duration = video_clip.duration
                    with AudioFileClip(audio_filepath) as audio_clip:
                        audio_duration = audio_clip.duration
                    
                    if abs(video_duration - audio_duration) < 1.0:
                        audio_exists = True
                        print(f"Found existing audio file '{audio_filepath}' with matching duration. Skipping extraction.")
                except Exception as check_err:
                    print(f"Could not verify existing audio file: {check_err}. Re-extracting...")

            if not audio_exists:
                try:
                    with VideoFileClip(filepath) as video_clip:
                        audio_clip = video_clip.audio
                        audio_clip.write_audiofile(audio_filepath)
                except Exception as e:
                    print(f"An error occurred extracting audio from video: {e}")
                    return

        with open(audio_filepath, "rb") as file:
            buffer_data = file.read()

        payload = {"buffer": buffer_data}

        options_dict = {
            "model": args.model,
            "smart_format": True,
            "utterances": True,
            "punctuate": True,
            "diarize": args.diarize,
        }
        if args.language:
            if args.language.lower() in {"auto", "detect"}:
                options_dict["detect_language"] = True
            else:
                options_dict["language"] = args.language
        options = PrerecordedOptions(**options_dict)

        print("Making request to deepgram")

        before = datetime.now()
        response = deepgram.listen.rest.v("1").transcribe_file(
            payload, options, timeout=httpx.Timeout(30000.0, connect=10.0)
        )
        after = datetime.now()
        print("Got response from deepgram")

        print(response.to_json(indent=4))

        # Check if the transcription contains words to avoid IndexError on silent audio files
        has_words = False
        try:
            if hasattr(response, 'results') and response.results:
                if response.results.channels and response.results.channels[0].alternatives:
                    if response.results.channels[0].alternatives[0].words:
                        has_words = True
        except Exception:
            pass

        if not has_words:
            print("No speech or words detected in the audio file. Generating empty subtitle file.")
            captions = ""
        else:
            transcription = DeepgramConverter(response)
            captions = srt(transcription)
        
        original_srt_path = f"{filepath}-captions.srt"
        cleaned_captions = cleanup_srt_punctuation(captions)
        with open(original_srt_path, "a", encoding="utf-8") as f:
            f.write(cleaned_captions)

        if args.translate_to:
            print(f"Translating subtitles to {args.translate_to} using DeepL...")
            deepl_api_key = os.getenv("DEEPL_API_KEY") or os.getenv("DEEPL_AUTH_KEY")
            if not deepl_api_key:
                print("Error: DEEPL_API_KEY or DEEPL_AUTH_KEY environment variable is not set.")
                print("Please set it in your environment or add it to your .env file to use translation.")
            else:
                try:
                    target_lang = args.translate_to.upper()
                    # DeepL-specific target language code overrides
                    if target_lang == "EN":
                        target_lang = "EN-US"
                    elif target_lang == "PT":
                        target_lang = "PT-BR"
                    
                    translated_srt_path = f"{filepath}-captions.{args.translate_to.lower()}.srt"
                    
                    # Translate and post-process
                    translated_content = translate_srt_content(cleaned_captions, deepl_api_key, target_lang)
                    cleaned_content = cleanup_srt_punctuation(translated_content)
                    with open(translated_srt_path, "w", encoding="utf-8") as f:
                        f.write(cleaned_content)
                        
                    print(f"Successfully translated subtitles. Saved to: {translated_srt_path}")
                except Exception as translate_err:
                    print(f"An error occurred during translation: {translate_err}")

        if should_remove_audio:
            os.remove(audio_filepath)

    except Exception as e:
        print(f"Exception: {e}")

if __name__ == "__main__":
    main()