import gradio as gr
import os
import torch, time
from huggingface_hub import hf_hub_download, snapshot_download, login
import sys
import gc
import base64
import torch

login(token=os.environ["HF_TOKEN"])
repo_id = os.environ["REPO_ID"]

torch.set_num_threads(1)
if torch.cuda.is_available():
    torch.backends.cudnn.benchmark = True

try:
    generate_file = hf_hub_download(repo_id=repo_id, filename="gen1.py", token=os.environ["HF_TOKEN"])
    os.system(f"cp {generate_file} ./gen1.py")
except Exception as e:
    print(f"Error downloading files: {e}")

from gen1 import setup_translation, translate_text as gen_translate

LOGO_PATH = "static/logo.png"
if os.path.isfile(LOGO_PATH):
    with open(LOGO_PATH, "rb") as f:
        LOGO_B64 = base64.b64encode(f.read()).decode()
    LOGO_HTML = f'<img src="data:image/png;base64,{LOGO_B64}" alt="Maruth Labs Logo" style="height:40px;">'
    FAVICON_HTML = f'<link rel="icon" type="image/png" href="data:image/png;base64,{LOGO_B64}">'
else:
    LOGO_HTML = '<div style="width:40px;height:40px;background:#ccc;border-radius:4px;"></div>'
    FAVICON_HTML = ''

def init_translation_model():
    success = setup_translation(repo_id=os.environ["REPO_ID"], token=os.environ["HF_TOKEN"])
    if success:
        print("Model loaded successfully!")
    else:
        print("Failed to load model")

def translate_text(source_text, source_lang, target_lang, temperature, top_k, repetition_penalty, max_tokens):
    return gen_translate(source_text, source_lang, target_lang, temperature, top_k, repetition_penalty, max_tokens)

languages = ["English", "Hindi", "Bengali", "Tamil", "Telugu", "Kannada", "Panjabi"]

theme_lock_css = """
.gradio-container .theme-toggle, .gradio-container button[aria-label*="theme"], .gradio-container button[title*="theme"],
.gradio-container .settings button, .gradio-container [data-testid="theme-toggle"] { display: none !important; }
:root { color-scheme: dark !important; }
body, .gradio-container {
    background-color: #0a1628 !important;
    color: #e6eef8 !important;
}

/* Mobile-responsive header fixes */
.main-header { padding: 10px 5px !important; }

.header-container {
    display: flex;
    align-items: center;
    justify-content: space-between;
    width: 100%;
    flex-wrap: wrap;
    gap: 10px;
}

.logo-section {
    display: flex;
    align-items: center;
    flex-shrink: 0;
    min-width: 0;
}

.logo-section h3 {
    margin-left: 8px;
    margin-top: 0;
    margin-bottom: 0;
    white-space: nowrap;
    font-size: 1rem;
}

.main-title {
    flex: 1;
    text-align: center;
    min-width: 0;
}

.main-title h1 {
    margin: 0;
    font-size: 1.5rem;
    line-height: 1.2;
    word-break: break-word;
}

/* Mobile responsiveness */
@media (max-width: 768px) {
    .header-container {
        flex-direction: column;
        align-items: center;
        text-align: center;
        gap: 15px;
    }
    
    .logo-section { order: 1; }
    
    .main-title {
        order: 2;
        width: 100%;
    }
    
    .main-title h1 {
        font-size: 1.3rem;
        margin: 0;
    }
    
    .logo-section h3 { font-size: 0.9rem; }
}

@media (max-width: 480px) {
    .main-title h1 { font-size: 1.1rem; }
    .logo-section h3 { font-size: 0.8rem; }
    .main-header { padding: 8px 3px !important; }
}
"""

locked_theme = gr.themes.Monochrome(primary_hue="blue", secondary_hue="slate", neutral_hue="slate").set(
    background_fill_primary="#0a1628",
    background_fill_secondary="#1f2937", 
    block_background_fill="#374151",
    border_color_primary="#374151",
    color_accent_soft="#2563eb",
    block_title_text_color="#e6eef8",
    block_label_text_color="#e6eef8",
    body_text_color="#e6eef8"
)

with gr.Blocks(
    title="Madhuram Translation - MaruthLabs",
    css=theme_lock_css,
    theme=locked_theme,
    js=f"""
    function() {{
        {f'document.head.insertAdjacentHTML("beforeend", `{FAVICON_HTML}`);' if FAVICON_HTML else ''}
        
        document.documentElement.setAttribute('data-theme', 'dark');
        document.body.classList.add('dark');
        document.body.classList.remove('light');
        const observer = new MutationObserver(function(mutations) {{
            mutations.forEach(function(mutation) {{
                mutation.addedNodes.forEach(function(node) {{
                    if (node.nodeType === 1) {{
                        const toggles = node.querySelectorAll('.theme-toggle, button[aria-label*="theme"], button[title*="theme"]');
                        toggles.forEach(toggle => toggle.style.display = 'none');
                    }}
                }});
            }});
        }});
        observer.observe(document.body, {{ childList: true, subtree: true }});
    }}
    """
) as demo:

    with gr.Row(elem_classes="main-header"):
        with gr.Column():
            gr.HTML(f"""
            <div class="header-container">
                <div class="logo-section">
                    {LOGO_HTML}
                    <h3>Maruth Labs</h3>
                </div>
                <div class="main-title">
                    <h1>Madhuram Translation Model</h1>
                </div>
            </div>
            """)

    with gr.Row():
        with gr.Column():
            gr.HTML("""
            <div style="background-color: #1e3a8a; border: 1px solid #3b82f6; border-radius: 8px; padding: 12px; margin: 10px 0;">
                <div style="display: flex; align-items: center; gap: 8px; color: #93c5fd;">
                    <svg width="16" height="16" viewBox="0 0 24 24" fill="currentColor">
                        <path d="M12 2C6.48 2 2 6.48 2 12s4.48 10 10 10 10-4.48 10-10S17.52 2 12 2zm-2 15l-5-5 1.41-1.41L10 14.17l7.59-7.59L19 8l-9 9z"/>
                    </svg>
                    <strong>Important:</strong> Use the original script of the languages for best translation results.
                </div>
            </div>
            """)

    with gr.Row(equal_height=False):
        with gr.Column(scale=1.5, elem_classes="settings-panel"):
            gr.Markdown("## Translation Settings")
            with gr.Row():
                source_lang = gr.Dropdown(choices=languages, label="Source Language", value="English")
                target_lang = gr.Dropdown(choices=languages, label="Target Language", value="Hindi")
            swap_btn = gr.Button("Swap Languages", variant="secondary", size="sm")

            with gr.Accordion("Advanced Settings", open=False):
                temperature = gr.Slider(0.001, 1.001, 0.001, step=0.1, label="Temperature")
                top_k = gr.Slider(1, 100, 10, step=1, label="Top-k")
                repetition_penalty = gr.Slider(1.0, 2.0, 1.2, step=0.1, label="Repetition Penalty")
                max_tokens = gr.Slider(100, 2000, 400, step=50, label="Max Tokens")

        with gr.Column(scale=2, elem_classes="translation-card"):
            gr.Markdown("## Translation Interface")
            source_text = gr.Textbox(label="Enter text to translate", placeholder="Type or paste your text here", lines=6, max_lines=12)
            with gr.Row():
                translate_btn = gr.Button("Translate", variant="primary", size="lg")
                clear_btn = gr.Button("Clear All", variant="secondary", size="lg")
            translated_text = gr.Textbox(label="Translation Result", lines=6, max_lines=12, interactive=False, placeholder="Translation will appear here")

    with gr.Row():
        with gr.Column():
            gr.Markdown("### Quick Examples")
            gr.Examples(
                examples=[
                    ["Hello, how are you today?", "English", "Hindi"],
                    ["তুমি কোথায় যাচ্ছ?", "Bengali", "English"],
                    ["நீங்கள் எப்படி இருக்கிறீர்கள்?", "Tamil", "Telugu"], 
                    ["ನಿನ್ನ ಹೆಸರು ಏನು?", "Kannada", "English"],
                    ["ਸਤ ਸ੍ਰੀ ਅਕਾਲ", "Panjabi", "Hindi"],   
                ],
                inputs=[source_text, source_lang, target_lang],
            )
            gr.Markdown("*Disclaimer - This is a demo version of Madhuram-Translate. It may occasionally generate incorrect or incomplete responses. Please verify independently. The complete model will be available through our own playground where the missing features will be incorporated.*", elem_classes="disclaimer")


    def swap_languages(src, tgt):
        return tgt, src
    
    def clear_all():
        return "", ""

    swap_btn.click(fn=swap_languages, inputs=[source_lang, target_lang], outputs=[source_lang, target_lang])
    clear_btn.click(fn=clear_all, outputs=[source_text, translated_text])    
    translate_btn.click(fn=translate_text, inputs=[source_text, source_lang, target_lang, temperature, top_k, repetition_penalty, max_tokens], outputs=[translated_text])
    demo.load(fn=init_translation_model)

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