#!/usr/bin/env python3
import os
import sys
import argparse
import subprocess
from dotenv import load_dotenv

# Load config
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)
else:
    load_dotenv()

# Add ai_transcriber to path so we can import modules
sys.path.append(os.path.join(script_dir, 'ai_transcriber'))

from ai_transcriber.translator import translate_srt
from ai_transcriber.utils import validate_and_repair_srt

def re_embed_subtitles(video_path, srt_path, output_path=None):
    """
    Re-embeds subtitles into an EXISTING video file, replacing the old tracks.
    Similar to embed_subtitles but we output to a temp file then overwrite.
    """
    if not os.path.exists(video_path) or not os.path.exists(srt_path):
        print("Error: Video or SRT file not found.")
        return False

    temp_output = video_path + ".temp.mp4"

    print(f"Re-embedding subtitles into: {video_path}...")
    
    # We map 0:v and 0:a from source, but ignore existing subs (0:s?)
    # Actually, simpler is to just map 0:v, 0:a and 1:0 (new sub)
    # and tell ffmpeg to make the new one default.
    # If the old file has subs, they might be kept unless we use -map -0:s (negative mapping)
    
    sub_codec = "mov_text" if video_path.lower().endswith(".mp4") else "srt"

    command = [
        "ffmpeg",
        "-i", video_path,
        "-i", srt_path,
        "-map", "0:v",
        "-map", "0:a",
        "-map", "1:0",      # The new subtitle
        "-c", "copy",       # Copy AV
        "-c:s", sub_codec,  # Encode Sub
        "-disposition:s:0", "default",
        "-metadata:s:s:0", "language=eng",
        "-metadata:s:s:0", "title=English (AI Translated)",
        "-y",
        "-v", "error",
        temp_output
    ]

    try:
        subprocess.run(command, check=True)
        # Move temp over original
        os.replace(temp_output, video_path)
        print(f"✅ Fixed: {video_path}")
        return True
    except subprocess.CalledProcessError as e:
        print(f"Error re-embedding: {e}")
        if os.path.exists(temp_output):
            os.remove(temp_output)
        return False

def process_recovery(folder_path, target_lang="English"):
    print(f"Scanning {folder_path} for incomplete translations...")
    
    count_fixed = 0
    count_skipped = 0

    video_extensions = ('.mp4', '.mkv', '.mov', '.avi')

    for root, dirs, files in os.walk(folder_path):
        for file in files:
            # We are looking for the SOURCE SRT files primarily
            if file.endswith(".srt") and not file.endswith(f".{target_lang}.srt"):
                source_srt_path = os.path.join(root, file)
                base_name = os.path.splitext(file)[0] # e.g. "movie" or "movie.subbed"
                
                # Check if this is a "subbed" file's srt? 
                # Usually source srt is "movie.srt". 
                # Video might be "movie.subbed.mp4" (if original was deleted)
                # OR "movie.mp4" (if not deleted)
                
                # Check if Translation exists
                expected_translation = os.path.join(root, f"{base_name}.{target_lang}.srt")
                
                if os.path.exists(expected_translation):
                    # It's already translated, nothing to do unless we force re-embed
                    continue
                
                print(f"\nFound untranslated transcript: {file}")
                
                # Try to translate
                with open(source_srt_path, "r", encoding="utf-8") as f:
                    content = f.read()
                
                new_srt_content = translate_srt(content, target_language=target_lang)
                
                if not new_srt_content:
                    print("❌ Translation failed (API Limit?). Skipping.")
                    continue
                
                # Save translation
                with open(expected_translation, "w", encoding="utf-8") as f:
                    f.write(new_srt_content)
                validate_and_repair_srt(expected_translation)
                
                # Now, find the video file to update
                # Case 1: Original name
                video_candidates = [
                    os.path.join(root, base_name + ".mp4"),
                    os.path.join(root, base_name + ".mkv"),
                    # Case 2: .subbed name (if original deleted)
                    os.path.join(root, base_name + ".subbed.mp4"), 
                ]
                
                found_video = None
                for v in video_candidates:
                    if os.path.exists(v):
                        found_video = v
                        break
                
                if found_video:
                    print(f"Found video to fix: {found_video}")
                    if re_embed_subtitles(found_video, expected_translation):
                        count_fixed += 1
                else:
                    print("Warning: Could not find a corresponding video file to embed into.")

    print(f"\nRecovery Complete. Fixed {count_fixed} files.")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: ./recover_and_fix.py <folder_path> [target_lang]")
        sys.exit(1)
        
    folder = sys.argv[1]
    lang = sys.argv[2] if len(sys.argv) > 2 else "English"
    
    process_recovery(folder, lang)