import os
import datetime
import gradio as gr
from groq import Groq
import torch
from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer

# ✅ Set API Key
os.environ["GROQ_API_KEY"] = "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"    # Replace with your API key
client = Groq(api_key=os.environ["GROQ_API_KEY"])

# ✅ Load M2M-100 Model
model_name = "facebook/m2m100_418M"
tokenizer = M2M100Tokenizer.from_pretrained(model_name)
model = M2M100ForConditionalGeneration.from_pretrained(model_name)

# ✅ Function to Generate English Script
def generate_script(topic, duration):
    try:
        words_per_minute = 130
        target_words = duration * words_per_minute

        response = client.chat.completions.create(
            messages=[{"role": "user", "content": f"Generate a plain, well-structured educational script about {topic} in English with approximately {target_words} words."}],
            model="llama-3.3-70b-versatile"
        )

        return response.choices[0].message.content
    except Exception as e:
        print(f"❌ Error in script generation: {str(e)}")
        return ""

# ✅ Function to Translate to Urdu
def translate_to_urdu(english_script):
    try:
        tokenizer.src_lang = "en"
        max_length = 500
        input_chunks = [english_script[i:i+max_length] for i in range(0, len(english_script), max_length)]

        translated_chunks = []
        for chunk in input_chunks:
            inputs = tokenizer(chunk, return_tensors="pt", truncation=True, max_length=1024).to(model.device)
            translated_tokens = model.generate(
                **inputs,
                max_length=1024,
                no_repeat_ngram_size=2,
                forced_bos_token_id=tokenizer.get_lang_id("ur"),
                num_beams=2
            )
            translated_chunks.append(tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0])

        return " ".join(translated_chunks)
    except Exception as e:
        return f"❌ Error in translation: {str(e)}"

# ✅ Save Edited Urdu Script
def save_edited_urdu(edited_urdu_script, topic):
    timestamp = datetime.datetime.now().strftime("%Y_%m_%d")
    filename = f"{topic}_Urdu_Final_{timestamp}.txt"
    filepath = os.path.join(os.getcwd(), filename)

    with open(filepath, "w", encoding="utf-8") as file:
        file.write(edited_urdu_script)

    return filepath

# ✅ Save File Utility
def save_file(content, filename):
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"{filename}_{timestamp}.txt"
    filepath = os.path.join(os.getcwd(), filename)

    with open(filepath, "w", encoding="utf-8") as file:
        file.write(content)

    return filepath

# ✅ Process Request
def process_request(topic, duration):
    eng_script = generate_script(topic, duration)
    return eng_script, "", ""  # Clears Urdu & Edited Urdu textboxes

def process_translation(eng_script):
    return translate_to_urdu(eng_script)

# ✅ Download Functions
def download_english(eng_script, topic):
    return save_file(eng_script, f"{topic}_Eng")

def download_urdu(urdu_script, topic):
    return save_file(urdu_script, f"{topic}_Urdu")

def download_final_urdu(edited_urdu_script, topic):
    return save_file(edited_urdu_script, f"{topic}_Urdu_Final")

# ✅ Gradio UI
with gr.Blocks() as app:
    gr.Markdown("# 🎬 AI-Powered Script Generator")

    topic_input = gr.Textbox(label="Enter Topic")
    duration_input = gr.Slider(minimum=1, maximum=30, step=1, label="Duration (minutes)")
    generate_button = gr.Button("Generate English Script")

    eng_output = gr.Textbox(label="Generated English Script", interactive=False)
    urdu_output = gr.Textbox(label="Translated Urdu Script", interactive=True)
    editable_urdu = gr.Textbox(label="Edited Urdu Script", interactive=True)

    generate_button.click(
        process_request,
        inputs=[topic_input, duration_input],
        outputs=[eng_output, urdu_output, editable_urdu]  # Clears Urdu & Edited Urdu textboxes
    )

    translate_button = gr.Button("Generate Urdu Script")
    translate_button.click(process_translation, inputs=[eng_output], outputs=[urdu_output])

    # ✅ Update editable Urdu box after translation
    translate_button.click(lambda x: x, inputs=[urdu_output], outputs=[editable_urdu])

    save_button = gr.Button("Save Edited Urdu Script")
    download_eng_button = gr.Button("📥 Download English Script")
    download_eng_file = gr.File()

    download_urdu_button = gr.Button("📥 Download Translated Urdu Script")
    download_urdu_file = gr.File()

    download_final_urdu_button = gr.Button("📥 Download Edited Urdu Script")
    download_final_urdu_file = gr.File()

    save_button.click(save_edited_urdu, inputs=[editable_urdu, topic_input], outputs=[download_final_urdu_file])
    download_eng_button.click(download_english, inputs=[eng_output, topic_input], outputs=[download_eng_file])
    download_urdu_button.click(download_urdu, inputs=[urdu_output, topic_input], outputs=[download_urdu_file])
    download_final_urdu_button.click(download_final_urdu, inputs=[editable_urdu, topic_input], outputs=[download_final_urdu_file])

app.launch()