import argparse
import os
import sys
from dotenv import load_dotenv

# Load environment variables from central .env_files directory
# Path: .../personal_development/video_transcription/ai_transcriber/main.py
# Target: .../personal_development/.env_files/.env.aitranscribe
script_dir = os.path.dirname(os.path.abspath(__file__))
env_path = os.path.abspath(os.path.join(script_dir, '../../.env_files/.env.aitranscribe'))

if os.path.exists(env_path):
    load_dotenv(env_path)
    # print(f"Loaded configuration from: {env_path}") # Optional: Uncomment for debugging
else:
    # Fallback: check local .env
    local_env = os.path.join(script_dir, '.env')
    if os.path.exists(local_env):
        load_dotenv(local_env)
    else:
        # Last resort: just try loading generic (cwd)
        load_dotenv()

from extractor import extract_audio, embed_subtitles
from transcriber import transcribe_audio, save_as_srt
from translator import translate_srt
from utils import validate_and_repair_srt
from diarizer import diarize_audio, merge_diarization_with_transcript

def save_srt_with_speakers(segments, output_path):
    """Helper to save SRT with speaker labels prepended to text."""
    def format_timestamp(seconds: float):
        whole_seconds = int(seconds)
        milliseconds = int((seconds - whole_seconds) * 1000)
        hours = whole_seconds // 3600
        minutes = (whole_seconds % 3600) // 60
        seconds = whole_seconds % 60
        return f"{hours:02d}:{minutes:02d}:{seconds:02d},{milliseconds:03d}"

    with open(output_path, "w", encoding="utf-8") as f:
        for i, segment in enumerate(segments, start=1):
            start = format_timestamp(segment["start"])
            end = format_timestamp(segment["end"])
            text = segment["text"].strip()
            speaker = segment.get("speaker", "")
            
            # Prepend speaker if present and not "Unknown"
            if speaker and speaker != "Unknown":
                text = f"[{speaker}]: {text}"
            
            f.write(f"{i}\n")
            f.write(f"{start} --> {end}\n")
            f.write(f"{text}\n\n")
    print(f"SRT saved to: {output_path}")

def process_file(file_path, args, source_lang=None):
    print(f"\n=== Processing: {file_path} ===")
    
    # 1. Extract Audio
    audio_path = extract_audio(file_path)
    
    # 2. Transcribe (Generate SRT)
    transcript_file = os.path.splitext(file_path)[0] + ".srt"
    transcript_exists = os.path.exists(transcript_file) and not args.force
    
    # Variable to hold final SRT path for embedding
    final_srt_path = transcript_file

    if transcript_exists:
        print(f"Transcript exists: {transcript_file}. Skipping transcription.")
        with open(transcript_file, "r", encoding="utf-8") as f:
            srt_content = f.read()
    else:
        # Transcribe
        result = transcribe_audio(audio_path, model_size=args.model, language=source_lang)
        segments = result["segments"]

        # Optional: Diarization
        if args.diarize:
            hf_token = args.hf_token or os.getenv("HF_TOKEN")
            if hf_token:
                print("Running Speaker Diarization...")
                diar_segments = diarize_audio(audio_path, hf_token=hf_token)
                if diar_segments:
                    segments = merge_diarization_with_transcript(segments, diar_segments)
                    print("Diarization merged into transcript.")
            else:
                print("Warning: --diarize requested but HF_TOKEN not provided. Skipping.")

        # Save SRT
        # Use simple save if no speakers, or custom if speakers
        if args.diarize:
             save_srt_with_speakers(segments, transcript_file)
        else:
             save_as_srt(result, transcript_file)
        
        # Validation
        validate_and_repair_srt(transcript_file)
        
        with open(transcript_file, "r", encoding="utf-8") as f:
            srt_content = f.read()

    # 3. Translate (Generate Translated SRT)
    translated_file = os.path.splitext(file_path)[0] + f".{args.lang}.srt"
    
    translation_success = False

    if os.path.exists(translated_file) and not args.force:
         print(f"Translation exists: {translated_file}. Skipping translation.")
         final_srt_path = translated_file
         translation_success = True
    else:
        # Only translate if there is content
        if srt_content:
            translated_srt_content = translate_srt(srt_content, target_language=args.lang)
            if translated_srt_content:
                with open(translated_file, "w", encoding="utf-8") as f:
                    f.write(translated_srt_content)
                print(f"Translation saved to: {translated_file}")
                validate_and_repair_srt(translated_file)
                final_srt_path = translated_file
                translation_success = True
            else:
                print("⚠️  TRANSLATION FAILED.")
                translation_success = False

    # 4. Embed Subtitles
    # SAFETY: If translation was intended but failed, do NOT embed/delete to prevent
    # replacing the video with one containing only untranslated subtitles.
    should_embed = args.embed
    if args.embed and not translation_success:
        print("\n❌ SAFETY HALT: Translation failed. Skipping embedding and deletion to preserve original file.")
        should_embed = False

    if should_embed:
        embed_subtitles(file_path, final_srt_path)

        import argparse
        import os
        import sys
        from dotenv import load_dotenv
        
        # Load environment variables from central .env_files directory
        # Path: .../personal_development/video_transcription/ai_transcriber/main.py
        # Target: .../personal_development/.env_files/.env.aitranscribe
        script_dir = os.path.dirname(os.path.abspath(__file__))
        env_path = os.path.abspath(os.path.join(script_dir, '../../.env_files/.env.aitranscribe'))
        
        if os.path.exists(env_path):
            load_dotenv(env_path)
            print(f"Loaded configuration from: {env_path}")
        else:
            # Fallback: check local .env
            local_env = os.path.join(script_dir, '.env')
            if os.path.exists(local_env):
                load_dotenv(local_env)
            else:
                # Last resort: just try loading generic (cwd)
                load_dotenv()
        
        from extractor import extract_audio, embed_subtitles
        
        from transcriber import transcribe_audio, save_as_srt
        
        from translator import translate_srt, translate_fallback_free
        
        from utils import validate_and_repair_srt
        
        from diarizer import diarize_audio, merge_diarization_with_transcript
        
        import tracker
        
        from tracker import JobStatus
        
        
        
        # ... (save_srt_with_speakers remains same)
        
        
        
        def process_file(file_path, args, source_lang=None):
        
            tracker.logger.info(f"=== Processing: {file_path} ===")
        
            
        
            # ... (Job init remains same) ...
        
            # ... (Step 1 Extract remains same) ...
        
            # ... (Step 2 Transcribe remains same) ...
        
        
        
                # 3. Translate (Generate Translated SRT)
        
                tracker.update_step(file_path, "step_translate", "processing")
        
                
        
                # Define paths
        
                base_translated = os.path.splitext(file_path)[0] + f".{args.lang}.srt"
        
                deep_translated = os.path.splitext(file_path)[0] + f".{args.lang}.deep_translate.srt"
        
                
        
                translated_file = base_translated # Default
        
                
        
                translation_success = False
        
                method_used = "None"
        
        
        
                if (os.path.exists(base_translated) or os.path.exists(deep_translated)) and not args.force:
        
                     if os.path.exists(deep_translated):
        
                         translated_file = deep_translated
        
                         method_used = "DeepTranslate (Existing)"
        
                     else:
        
                         method_used = "Gemini (Existing)"
        
                         
        
                     tracker.logger.info(f"Translation exists: {translated_file} ({method_used}). Skipping translation.")
        
                     final_srt_path = translated_file
        
                     translation_success = True
        
                else:
        
                    # Only translate if there is content
        
                    if srt_content:
        
                        # Attempt 1: Gemini
        
                        translated_srt_content = translate_srt(srt_content, target_language=args.lang)
        
                        
        
                        if translated_srt_content:
        
                            with open(base_translated, "w", encoding="utf-8") as f:
        
                                f.write(translated_srt_content)
        
                            tracker.logger.info(f"Translation saved to: {base_translated} (Gemini)")
        
                            validate_and_repair_srt(base_translated)
        
                            final_srt_path = base_translated
        
                            translation_success = True
        
                            method_used = "Gemini"
        
                        else:
        
                            # Attempt 2: Fallback
        
                            tracker.logger.warning("Gemini translation failed. Attempting Free Fallback...")
        
                            
        
                            lang_map = {
        
                                "English": "en", "French": "fr", "Spanish": "es", 
        
                                "German": "de", "Italian": "it", "Portuguese": "pt",
        
                                "Russian": "ru", "Japanese": "ja", "Chinese": "zh-CN"
        
                            }
        
                            target_code = lang_map.get(args.lang, "en")
        
                            
        
                            translated_srt_content = translate_fallback_free(srt_content, target_language=target_code)
        
                            
        
                            if translated_srt_content:
        
                                translated_file = deep_translated
        
                                with open(translated_file, "w", encoding="utf-8") as f:
        
                                    f.write(translated_srt_content)
        
                                tracker.logger.info(f"Translation saved to: {translated_file} (DeepTranslate)")
        
                                validate_and_repair_srt(translated_file)
        
                                final_srt_path = translated_file
        
                                translation_success = True
        
                                method_used = "DeepTranslate"
        
                            else:
        
                                tracker.logger.error("TRANSLATION FAILED (Both Gemini and Fallback).")
        
                                tracker.update_step(file_path, "step_translate", "failed")
        
                                translation_success = False
        
        
        
                if translation_success:
        
                     tracker.update_step(file_path, "step_translate", "done")
        
                     tracker.logger.info(f"Translation Method: {method_used}")
        
        
        
                # 4. Embed Subtitles
        
        
                tracker.update_step(file_path, "step_embed", "processing")
                should_embed = args.embed
                if args.embed and not translation_success:
                    tracker.logger.warning("SAFETY HALT: Translation failed. Skipping embedding and deletion to preserve original file.")
                    should_embed = False
        
                if should_embed:
                    embed_subtitles(file_path, final_srt_path)
        
                    # 5. Delete Source File (Optional & Risky)
                    if args.delete_source:
                        if args.embed:
                            # Safety: Ensure the new subbed video exists before deleting the old one
                            base, ext = os.path.splitext(file_path)
                            expected_output = f"{base}.subbed{ext}"
                            
                            if os.path.exists(expected_output):
                                try:
                                    os.remove(file_path)
                                    tracker.logger.info(f"SOURCE DELETED: Original file '{file_path}' has been removed.")
                                except OSError as e:
                                    tracker.logger.error(f"Error: Could not delete source file: {e}")
                            else:
                                tracker.logger.error(f"SAFETY ABORT: Source file NOT deleted. Could not find expected output '{expected_output}'.")
                        else:
                            tracker.logger.warning("SAFETY ABORT: Source file NOT deleted. You must enable --embed to safely replace the video.")
                tracker.update_step(file_path, "step_embed", "done")
        
                # 5. Cleanup Audio
                if args.cleanup:
                    try:
                        os.remove(audio_path)
                        tracker.logger.info(f"Cleanup: Removed temporary audio file {audio_path}")
                    except OSError as e:
                        tracker.logger.warning(f"Warning: Could not remove audio file: {e}")
        
                # Mark Complete
                if translation_success:
                    tracker.update_job_status(file_path, JobStatus.COMPLETED)
                else:
                    # If translation failed but we didn't crash, we technically finished the run but result is partial
                    tracker.update_job_status(file_path, JobStatus.FAILED, error="Translation failed")
        
            except Exception as e:
                tracker.logger.exception(f"Job Failed for {file_path}")
                tracker.update_job_status(file_path, JobStatus.FAILED, error=str(e))
                # Don't exit, allow other files to process
                return
        
        def main():
            parser = argparse.ArgumentParser(description="AI Video Transcriber & Translator")
            parser.add_argument("input", nargs='?', help="Path to video file or directory")
            parser.add_argument("--model", default="auto", choices=["auto", "tiny", "base", "small", "medium", "large"], help="Whisper model size (default: auto)")
            parser.add_argument("--lang", default="English", help="Target language for translation (default: English)")
            parser.add_argument("--source-lang", help="Source language of the audio (e.g., 'fr', 'es'). If omitted, you will be prompted.")
            parser.add_argument("--force", action="store_true", help="Overwrite existing transcript/translation files")
            
            # New Arguments
            parser.add_argument("--cleanup", action="store_true", help="Delete the temporary .wav file after processing")
            parser.add_argument("--embed", action="store_true", help="Embed the final subtitles into the video (Soft Subs)")
            parser.add_argument("--diarize", action="store_true", help="Enable speaker diarization (requires HF_TOKEN)")
            parser.add_argument("--hf-token", help="HuggingFace Token for pyannote.audio (or set HF_TOKEN env var)")
            parser.add_argument("--delete-source", action="store_true", help="Delete the original video file AFTER successful embedding")
            parser.add_argument("--retry-failed", action="store_true", help="Retry only jobs marked as FAILED in the database")
            
            args = parser.parse_args()
        
            if not os.getenv("GEMINI_API_KEY"):
                print("Warning: GEMINI_API_KEY environment variable not set. Translation step will fail.")
            
            # Handling Retry Logic
            if args.retry_failed:
                print("Retrying failed jobs from database...")
                failed_files = tracker.get_failed_jobs()
                if not failed_files:
                    print("No failed jobs found.")
                    return
                
                # We need args.source_lang logic here too if needed, but for retries we might assume context
                # For simplicity, we'll prompt if missing just like normal run
                
                # Determine source language (Prompt if missing)
                source_lang = args.source_lang
                if not source_lang:
                    print("\n--- Audio Configuration ---")
                    user_input = input("Enter the source language of the video(s) (e.g., 'French', 'es').\nPress Enter to use Whisper's auto-detection: ").strip()
                    if user_input:
                        source_lang = user_input
                    else:
                        source_lang = None # Let Whisper auto-detect
                        print("Selected: Auto-detect")
        
                for file_path in failed_files:
                    if os.path.exists(file_path):
                        process_file(file_path, args, source_lang)
                    else:
                        print(f"Skipping missing file: {file_path}")
                return
        
            # Normal Logic
            if not args.input:
                parser.print_help()
                sys.exit(1)
        
            # Determine source language (Prompt if missing)
            source_lang = args.source_lang
            if not source_lang:
                print("\n--- Audio Configuration ---")
                user_input = input("Enter the source language of the video(s) (e.g., 'French', 'es').\nPress Enter to use Whisper's auto-detection: ").strip()
                if user_input:
                    source_lang = user_input
                else:
                    source_lang = None # Let Whisper auto-detect
                    print("Selected: Auto-detect")
            
            if os.path.isfile(args.input):
                process_file(args.input, args, source_lang)
            elif os.path.isdir(args.input):
                video_extensions = ('.mp4', '.mkv', '.mov', '.avi', '.webm', '.flv', '.wmv', '.m4v')
                found_files = False
                for root, dirs, files in os.walk(args.input):
                    for file in files:
                        if file.lower().endswith(video_extensions):
                            found_files = True
                            file_path = os.path.join(root, file)
                            process_file(file_path, args, source_lang)
                if not found_files:
                     print(f"No video files found in {args.input}")
            else:
                print(f"Error: Invalid input path '{args.input}'")
                sys.exit(1)
        
        if __name__ == "__main__":
            main()
def main():
    parser = argparse.ArgumentParser(description="AI Video Transcriber & Translator")
    parser.add_argument("input", help="Path to video file or directory")
    parser.add_argument("--model", default="auto", choices=["auto", "tiny", "base", "small", "medium", "large"], help="Whisper model size (default: auto)")
    parser.add_argument("--lang", default="English", help="Target language for translation (default: English)")
    parser.add_argument("--source-lang", help="Source language of the audio (e.g., 'fr', 'es'). If omitted, you will be prompted.")
    parser.add_argument("--force", action="store_true", help="Overwrite existing transcript/translation files")
    
    # New Arguments
    parser.add_argument("--cleanup", action="store_true", help="Delete the temporary .wav file after processing")
    parser.add_argument("--embed", action="store_true", help="Embed the final subtitles into the video (Soft Subs)")
    parser.add_argument("--diarize", action="store_true", help="Enable speaker diarization (requires HF_TOKEN)")
    parser.add_argument("--hf-token", help="HuggingFace Token for pyannote.audio (or set HF_TOKEN env var)")
    parser.add_argument("--delete-source", action="store_true", help="Delete the original video file AFTER successful embedding")
    
    args = parser.parse_args()

    if not os.getenv("GEMINI_API_KEY"):
        print("Warning: GEMINI_API_KEY environment variable not set. Translation step will fail.")

    # Determine source language (Prompt if missing)
    source_lang = args.source_lang
    if not source_lang:
        print("\n--- Audio Configuration ---")
        user_input = input("Enter the source language of the video(s) (e.g., 'French', 'es').\nPress Enter to use Whisper's auto-detection: ").strip()
        if user_input:
            source_lang = user_input
        else:
            source_lang = None # Let Whisper auto-detect
            print("Selected: Auto-detect")
    
    if os.path.isfile(args.input):
        process_file(args.input, args, source_lang)
    elif os.path.isdir(args.input):
        video_extensions = ('.mp4', '.mkv', '.mov', '.avi', '.webm', '.flv', '.wmv', '.m4v')
        found_files = False
        for root, dirs, files in os.walk(args.input):
            for file in files:
                if file.lower().endswith(video_extensions):
                    found_files = True
                    file_path = os.path.join(root, file)
                    process_file(file_path, args, source_lang)
        if not found_files:
             print(f"No video files found in {args.input}")
    else:
        print(f"Error: Invalid input path '{args.input}'")
        sys.exit(1)

if __name__ == "__main__":
    main()