import json
import asyncio
import textwrap
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import google.generativeai as genai
import edge_tts
from moviepy import VideoFileClip, AudioFileClip, ImageClip, CompositeVideoClip
from dotenv import load_dotenv
import os

load_dotenv()
raw_api_key = os.getenv("GENAI_API_KEY")

# 2. Check if it exists
if not raw_api_key:
    raise ValueError("GENAI_API_KEY is missing in Hugging Face Secrets!")

CLEAN_API_KEY = raw_api_key.strip()

# 4. Configure Gemini with the clean key
genai.configure(api_key=CLEAN_API_KEY)

# GENAI_API_KEY = os.getenv("GENAI_API_KEY")
# genai.configure(api_key=GENAI_API_KEY)

def format_timestamp(seconds):
    """Converts seconds (float) to WebVTT format (HH:MM:SS.mmm)"""
    milliseconds = int((seconds % 1) * 1000)
    minutes = int(seconds // 60)
    hours = int(minutes // 60)
    minutes = minutes % 60
    seconds = int(seconds % 60)
    return f"{hours:02}:{minutes:02}:{seconds:02}.{milliseconds:03}"

async def generate_dubbing(text, voice, output_file):
    communicate = edge_tts.Communicate(text, voice)
    await communicate.save(output_file)

def translate_and_dub(video_path, target_lang, gender="Female"):
    base_name = os.path.splitext(video_path)[0]
    audio_path = f"{base_name}_temp.mp3"
    dub_audio_path = f"{base_name}_dub.mp3"
    output_video_path = f"{base_name}_dubbed.mp4"
    output_sub_path = f"{base_name}_subs.vtt" # We create a VTT file now

    # 1. EXTRACT AUDIO & INFO
    with VideoFileClip(video_path) as video:
        video.audio.write_audiofile(audio_path, logger=None)
        duration = video.duration

    # 2. AI TRANSLATION & TIMESTAMPS
    model = genai.GenerativeModel("gemini-3-flash-preview")
    
    if target_lang == "Chinese":
        voice = "zh-CN-YunxiNeural" if gender == "Male" else "zh-CN-XiaoxiaoNeural"
        lang_prompt = "Simplified Chinese"
    else:
        voice = "vi-VN-NamMinhNeural" if gender == "Male" else "vi-VN-HoaiMyNeural"
        lang_prompt = "Vietnamese"

    prompt = f"""
    Listen to this audio. Return a JSON list of segments.
    For each segment, translate the spoken content into {lang_prompt}.
    Format:
    [
      {{"start": 0.0, "end": 2.5, "text": "Translated text here"}},
      {{"start": 2.5, "end": 5.0, "text": "Next text here"}}
    ]
    Use seconds for timestamps. 
    Ensure the segments cover the whole video.
    """

    print("Sending to Gemini...")
    audio_file = genai.upload_file(path=audio_path)
    response = model.generate_content([prompt, audio_file], generation_config={"response_mime_type": "application/json"})
    
    try:
        segments = json.loads(response.text)
    except json.JSONDecodeError:
        segments = [{"start": 0, "end": duration, "text": "Translation Error: Could not parse JSON."}]

    # 3. GENERATE VTT FILE (SUBTITLES)
    print("Creating subtitles...")
    vtt_content = "WEBVTT\n\n"
    full_text_for_dub = []
    
    for seg in segments:
        start_time = format_timestamp(float(seg['start']))
        end_time = format_timestamp(float(seg['end']))
        text = seg['text']
        
        # Add to VTT
        vtt_content += f"{start_time} --> {end_time}\n{text}\n\n"
        
        # Collect text for dubbing
        full_text_for_dub.append(text)

    # Save VTT file
    with open(output_sub_path, "w", encoding="utf-8") as f:
        f.write(vtt_content)

    # 4. GENERATE DUBBING AUDIO
    print("Generating voice...")
    full_text = " ".join(full_text_for_dub)
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.run_until_complete(generate_dubbing(full_text, voice, dub_audio_path))

    # 5. MERGE AUDIO ONLY (No video re-encoding needed usually, but MoviePy is safest)
    print("Merging new audio...")
    with VideoFileClip(video_path) as video:
        with AudioFileClip(dub_audio_path) as dub:
            # Handle duration mismatch
            if dub.duration > video.duration:
                dub = dub.subclipped(0, video.duration)
            
            final_clip = video.with_audio(dub)
            final_clip.write_videofile(output_video_path, codec="libx264", audio_codec="aac", logger=None)

    # Cleanup temp files
    if os.path.exists(audio_path): os.remove(audio_path)
    # Keeping dub audio and vtt might be useful for the user

    return output_video_path, output_sub_path, segments