import os
import torch
import spaces
import gradio as gr
from threading import Thread
from collections.abc import Iterator
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

MAX_MAX_NEW_TOKENS = 4096
MAX_INPUT_TOKEN_LENGTH = 4096
DEFAULT_MAX_NEW_TOKENS = 2048
HF_TOKEN = os.environ['HF_TOKEN']

model_id = "ai4bharat/IndicTrans3-beta"
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto", token=HF_TOKEN)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-3B-Instruct")

LANGUAGES = {
    "Hindi": "hin_Deva",
    "Bengali": "ben_Beng",
    "Telugu": "tel_Telu",
    "Marathi": "mar_Deva",
    "Tamil": "tam_Taml",
    "Urdu": "urd_Arab",
    "Gujarati": "guj_Gujr",
    "Kannada": "kan_Knda",
    "Odia": "ori_Orya",
    "Malayalam": "mal_Mlym",
    "Punjabi": "pan_Guru",
    "Assamese": "asm_Beng",
    "Maithili": "mai_Mith",
    "Santali": "sat_Olck",
    "Kashmiri": "kas_Arab",
    "Nepali": "nep_Deva",
    "Sindhi": "snd_Arab",
    "Konkani": "kok_Deva",
    "Dogri": "dgo_Deva",
    "Manipuri": "mni_Beng",
    "Bodo": "brx_Deva"
}

def format_message_for_translation(message, target_lang):
    return f"Translate the following text to {target_lang}: {message}"

@spaces.GPU
def translate_message(
    message: str,
    chat_history: list[dict],
    target_language: str = "Hindi",
    max_new_tokens: int = 1024,
    temperature: float = 0.6,
    top_p: float = 0.9,
    top_k: int = 50,
    repetition_penalty: float = 1.2,
) -> Iterator[str]:
    conversation = []
    
    translation_request = format_message_for_translation(message, target_language)
    print(f"Translation request: {translation_request}")
    conversation.append({"role": "user", "content": translation_request})

    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt", add_generation_prompt=True)
    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
        gr.Warning(f"Trimmed input as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
    input_ids = input_ids.to(model.device)

    streamer = TextIteratorStreamer(tokenizer, timeout=240.0, skip_prompt=True, skip_special_tokens=True)
    generate_kwargs = dict(
        {"input_ids": input_ids},
        streamer=streamer,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        top_p=top_p,
        top_k=top_k,
        temperature=temperature,
        num_beams=1,
        repetition_penalty=repetition_penalty,
    )
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    outputs = []
    for text in streamer:
        outputs.append(text)
        yield "".join(outputs)

def store_feedback(rating, feedback_text):
    if not rating:
        gr.Warning("Please select a rating before submitting feedback.", duration=5)
        return None
    
    if not feedback_text or feedback_text.strip() == "":
        gr.Warning("Please provide some feedback before submitting.", duration=5)
        return None
    
    gr.Info("Feedback submitted successfully!")
    return "Thank you for your feedback!"

css = """
# body {
#     background-color: #f7f7f7;
# }
.feedback-section {
    margin-top: 30px;
    border-top: 1px solid #ddd;
    padding-top: 20px;
}
.container {
    max-width: 90%;
    margin: 0 auto;
}
.language-selector {
    margin-bottom: 20px;
    padding: 10px;
    background-color: #ffffff;
    border-radius: 8px;
    box-shadow: 0 2px 5px rgba(0,0,0,0.1);
}
.advanced-options {
    margin-top: 20px;
}
"""

DESCRIPTION = """\
IndicTrans3 is the latest state-of-the-art (SOTA) translation model from AI4Bharat, designed to handle translations across <b>22 Indic languages</b> with high accuracy. It supports <b>document-level machine translation (MT)</b> and is built to match the performance of other leading SOTA models. <br>
📢 <b>Training data will be released soon!</b>  
<h3>🔹 Features</h3>  
✅ Supports <b>22 Indic languages</b>  
✅ Enables <b>document-level translation</b>  
✅ Achieves <b>SOTA performance</b> in Indic MT  
✅ Optimized for <b>real-world applications</b>  
<h3>🚀 Try It Out!</h3>  
1️⃣ Enter text in any supported language  
2️⃣ Select the target language  
3️⃣ Click <b>Translate</b> and get high-quality results!  
Built for <b>linguistic diversity and accessibility</b>, IndicTrans3 is a major step forward in <b>Indic language AI</b>.  
💡 <b>Source:</b> AI4Bharat | Powered by Hugging Face  
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_classes="container"):
        gr.Markdown("# 🌏 IndicTrans3-beta 🚀: Multilingual Translation for 22 Indic Languages </center>")
        gr.Markdown(DESCRIPTION)
        
        target_language = gr.Dropdown(
            list(LANGUAGES.keys()), 
            value="Hindi",
            label="Which language would you like to translate to?",
            elem_id="language-dropdown"
        )
    
        chatbot = gr.Chatbot(height=400, elem_id="chatbot")
        
        with gr.Row():
            msg = gr.Textbox(
                placeholder="Enter text to translate...",
                show_label=False,
                container=False,
                scale=9
            )
            submit_btn = gr.Button("Translate", scale=1)
        
        gr.Examples(
            examples=[
                "The Taj Mahal stands majestically along the banks of river Yamuna, a timeless symbol of eternal love.",
                "Kumbh Mela is the world's largest gathering of people, where millions of pilgrims bathe in sacred rivers for spiritual purification.",
                "India's classical dance forms like Bharatanatyam, Kathak, and Odissi beautifully blend rhythm, expression, and storytelling.",
                "Ayurveda, the ancient Indian medical system, focuses on holistic wellness through natural herbs and balanced living.",
                "During Diwali, homes across India are decorated with oil lamps, colorful rangoli patterns, and twinkling lights to celebrate the victory of light over darkness."
            ],
            inputs=msg
        )
        
        
        with gr.Accordion("Provide Feedback", open=True):
            gr.Markdown("## Rate Translation & Provide Feedback 📝")
            gr.Markdown("Help us improve the translation quality by providing your feedback.")
            with gr.Row():
                rating = gr.Radio(
                    ["1", "2", "3", "4", "5"],
                    label="Translation Rating (1-5)"
                )
            
            feedback_text = gr.Textbox(
                placeholder="Share your feedback about the translation...",
                label="Feedback",
                lines=3
            )
            
            feedback_submit = gr.Button("Submit Feedback")
            feedback_result = gr.Textbox(label="", visible=False)

        with gr.Accordion("Advanced Options", open=False, elem_classes="advanced-options"):
            max_new_tokens = gr.Slider(
                label="Max new tokens",
                minimum=1,
                maximum=MAX_MAX_NEW_TOKENS,
                step=1,
                value=DEFAULT_MAX_NEW_TOKENS,
            )
            temperature = gr.Slider(
                label="Temperature",
                minimum=0.1,
                maximum=1.0,
                step=0.1,
                value=0.1,
            )
            top_p = gr.Slider(
                label="Top-p (nucleus sampling)",
                minimum=0.05,
                maximum=1.0,
                step=0.05,
                value=0.9,
            )
            top_k = gr.Slider(
                label="Top-k",
                minimum=1,
                maximum=100,
                step=1,
                value=50,
            )
            repetition_penalty = gr.Slider(
                label="Repetition penalty",
                minimum=1.0,
                maximum=2.0,
                step=0.05,
                value=1.0,
            )
        
        chat_state = gr.State([])
        
        def user(user_message, history, target_lang):
            return "", history + [[user_message, None]]
        
        def bot(history, target_lang, max_tokens, temp, top_p_val, top_k_val, rep_penalty):
            user_message = history[-1][0]
            history[-1][1] = ""
            
            for chunk in translate_message(
                user_message, 
                history[:-1], 
                target_lang, 
                max_tokens, 
                temp, 
                top_p_val, 
                top_k_val, 
                rep_penalty
            ):
                history[-1][1] = chunk
                yield history
        
        msg.submit(
            user,
            [msg, chatbot, target_language],
            [msg, chatbot],
            queue=False
        ).then(
            bot,
            [chatbot, target_language, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
            chatbot
        )
        
        submit_btn.click(
            user,
            [msg, chatbot, target_language],
            [msg, chatbot],
            queue=False
        ).then(
            bot,
            [chatbot, target_language, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
            chatbot
        )
        
        feedback_submit.click(
            fn=store_feedback, 
            inputs=[rating, feedback_text], 
            outputs=feedback_result
        )

if __name__ == "__main__":
    demo.launch()