# Updated for Inachi-Core (Elephant AI) - Text & Audio Dual Mode
import os
import time
import gradio as gr
from src.mimo_audio.mimo_audio import MimoAudio

class InachiProEngine:
    def __init__(self, model_path, tokenizer_path):
        # MiMo-V2.5-Pro load කිරීම
        self.model = MimoAudio(model_path, tokenizer_path)

    def generate(self, text_input, audio_input, language_choice):
        # Audio හෝ Text යන දෙකෙන් ඕනෑම එකක් process කිරීමේ හැකියාව
        audio_tag = LANGUAGE_TAGS.get(language_choice, "")
        
        try:
            start = time.time()
            # මෙතනදී text_input එක කෙලින්ම model එකට pass කළ හැකියි
            # MiMo-Pro හි text-to-text හෝ audio-to-text functions පාවිච්චි වේ
            if audio_input:
                result = self.model.asr_sft(audio_input, audio_tag=audio_tag)
            else:
                # Text chat logic
                result = self.model.chat(text_input) 
            
            elapsed = time.time() - start
            return result, f"🚀 Processed in {elapsed:.2f}s"
        except Exception as e:
            return "", f"❌ Error: {str(e)}"

# UI එකට Textbox එකක් ඇතුළත් කිරීම
def create_dual_interface(engine):
    with gr.Blocks(theme=gr.themes.Default(primary_hue="blue")) as iface:
        gr.Markdown("# 🔱 INACHI-CORE | MiMo-V2.5-Pro")
        
        with gr.Row():
            with gr.Column(scale=1):
                audio_in = gr.Audio(label="Audio Input (Optional)", type="filepath")
                text_in = gr.Textbox(label="Message / Prompt", placeholder="Type your command here...")
                lang = gr.Radio(choices=["Auto", "Chinese", "English"], value="Auto", label="Language Context")
                submit_btn = gr.Button("Execute Command", variant="primary")
            
            with gr.Column(scale=1):
                chat_out = gr.Textbox(label="Inachi Response", lines=12)
                status = gr.Label(label="System Heartbeat")

        submit_btn.click(
            fn=engine.generate,
            inputs=[text_in, audio_in, lang],
            outputs=[chat_out, status]
        )
    return iface