import gradio as gr
from src.indus.engines.router.engine import router_engine
from src.indus.engines.model.adapter import model_engine

def predict(message, history):
    # Determine what the user sent (message can be a dict if multimodal=True)
    modalities = {}
    prompt = ""
    
    if isinstance(message, dict):
        prompt = message.get("text", "")
        files = message.get("files", [])
        if files:
            # Simple check for demo purposes
            for file in files:
                if any(ext in file.lower() for ext in ['.jpg', '.png', '.jpeg', '.webp']):
                    modalities['image'] = True
                elif any(ext in file.lower() for ext in ['.mp3', '.wav', '.ogg']):
                    modalities['audio'] = True
    else:
        prompt = message

    context = {
        "prompt": prompt,
        "modalities": modalities
    }
    
    # 1. Route the prompt
    routed_context = router_engine.classify(context)
    
    # 2. Extract Required Capabilities
    profile = routed_context.get("__capability_profile__")
    target_cap = profile.target_role if profile else "reasoning"
    
    # 3. Generate response using the local Indus Engine
    response = model_engine.generate(prompt, required_capabilities=[target_cap])
    
    return f"[Router: {target_cap}]\n\n{response}"

with gr.Blocks(title="Indus AI - Omni Interface") as demo:
    gr.Markdown("# 🚀 Indus AI: Omni Interface")
    gr.Markdown("Welcome to your local AI OS. This interface uses the `Multimodal Capability Router` to dynamically dispatch your input to the optimal trained model.")
    
    chatbot = gr.Chatbot()
    
    chat_interface = gr.ChatInterface(
        fn=predict,
        multimodal=True,
        chatbot=chatbot,
        textbox=gr.MultimodalTextbox(file_types=["image", "audio"], placeholder="Ask anything, or attach an image/audio...")
    )
    
    with gr.Accordion("Report a Failure (Data Flywheel)", open=False):
        gr.Markdown("Did the AI make a mistake? Provide the correct answer below. It will automatically be ingested into the next training run.")
        correction_input = gr.Textbox(label="Expected Correct Answer")
        submit_btn = gr.Button("Submit Failure to Flywheel")
        feedback_status = gr.Markdown("")
        
        def log_failure(history, correct_answer):
            if not history or len(history) == 0:
                return "No conversation history found to log."
                
            last_user_msg = history[-1][0]
            last_bot_msg = history[-1][1]
            
            # Extract text if multimodal
            if isinstance(last_user_msg, tuple):
                last_user_msg = last_user_msg[0]
                
            from src.indus.engines.evaluation.logger import failure_logger
            filepath = failure_logger.log_failure(
                prompt=str(last_user_msg),
                model_answer=str(last_bot_msg),
                expected_answer=correct_answer
            )
            return f"✅ Failure successfully logged to `{filepath}`! It will be included in your next dataset generation."
            
        submit_btn.click(
            fn=log_failure,
            inputs=[chatbot, correction_input],
            outputs=[feedback_status]
        )

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