import gradio as gr
import torch
import torchaudio
import numpy as np
from transformers import AutoProcessor, SeamlessM4Tv2Model
from datetime import datetime
import time

class JarvisTranslator:
    def __init__(self, model_name: str = "facebook/seamless-m4t-v2-large"):
        self.processor = AutoProcessor.from_pretrained(model_name)
        self.model = SeamlessM4Tv2Model.from_pretrained(model_name)
        self.sample_rate = self.model.config.sampling_rate
        
        self.language_codes = {
            "English": "eng",
            "Spanish": "spa",
            "French": "fra",
            "German": "deu",
            "Italian": "ita",
            "Portuguese": "por",
            "Russian": "rus",
            "Chinese": "cmn",
            "Japanese": "jpn"
        }

    def translate(self, text: str, src_lang: str, tgt_lang: str) -> tuple[int, np.ndarray]:
        inputs = self.processor(text=text, src_lang=src_lang, return_tensors="pt")
        audio_array = self.model.generate(**inputs, tgt_lang=tgt_lang)[0].cpu().numpy().squeeze()
        return self.sample_rate, audio_array

def create_jarvis_interface():
    # Custom CSS for Jarvis-like theme
    css = """
    #jarvis-container {
        background-color: #000000;
        color: #00ffff;
        font-family: 'Courier New', monospace;
        padding: 20px;
        border-radius: 10px;
        border: 2px solid #00ffff;
    }
    
    #status-circle {
        width: 150px;
        height: 150px;
        border: 4px solid #00ffff;
        border-radius: 50%;
        margin: 20px auto;
        position: relative;
        animation: pulse 2s infinite;
    }
    
    @keyframes pulse {
        0% { box-shadow: 0 0 0 0 rgba(0, 255, 255, 0.4); }
        70% { box-shadow: 0 0 0 20px rgba(0, 255, 255, 0); }
        100% { box-shadow: 0 0 0 0 rgba(0, 255, 255, 0); }
    }
    
    .custom-button {
        background-color: transparent !important;
        border: 2px solid #00ffff !important;
        color: #00ffff !important;
        font-family: 'Courier New', monospace !important;
    }
    
    .custom-button:hover {
        background-color: rgba(0, 255, 255, 0.1) !important;
    }
    
    .status-text {
        color: #00ffff;
        text-align: center;
        font-size: 1.2em;
        margin: 10px 0;
    }
    
    .time-display {
        position: absolute;
        top: 10px;
        right: 10px;
        color: #00ffff;
        font-family: 'Courier New', monospace;
    }
    """

    translator = JarvisTranslator()

    def update_status():
        return f"JARVIS AI SYSTEM ACTIVE\nTime: {datetime.now().strftime('%H:%M:%S')}"

    def process_command(text, src_lang, tgt_lang):
        status = f"Processing command: {text}\nSource: {src_lang} → Target: {tgt_lang}"
        time.sleep(1)  # Simulate processing
        try:
            sample_rate, audio = translator.translate(text, 
                                                    translator.language_codes[src_lang],
                                                    translator.language_codes[tgt_lang])
            return audio, status + "\nStatus: Translation complete"
        except Exception as e:
            return None, f"Error: {str(e)}"

    with gr.Blocks(css=css, title="JARVIS AI") as demo:
        with gr.Column(elem_id="jarvis-container"):
            gr.Markdown("# JARVIS AI TRANSLATION SYSTEM")
            
            # Status display
            status_html = gr.HTML(value="<div id='status-circle'></div>", show_label=False)
            status_text = gr.Textbox(label="System Status", value=update_status)
            
            with gr.Row():
                text_input = gr.Textbox(
                    label="Command Input",
                    placeholder="Enter text to translate...",
                    lines=3
                )
                
            with gr.Row():
                src_lang = gr.Dropdown(
                    choices=list(translator.language_codes.keys()),
                    value="English",
                    label="Source Language"
                )
                tgt_lang = gr.Dropdown(
                    choices=list(translator.language_codes.keys()),
                    value="Spanish",
                    label="Target Language"
                )
                
            with gr.Row():
                process_btn = gr.Button("Execute Translation", elem_classes=["custom-button"])
                
            audio_output = gr.Audio(
                label="Translated Output",
                type="numpy"
            )
            
            # Event handlers
            process_btn.click(
                fn=process_command,
                inputs=[text_input, src_lang, tgt_lang],
                outputs=[audio_output, status_text]
            )
            
            demo.load(
                fn=update_status,
                outputs=status_text,
                every=1  # Update every second
            )

    return demo

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