import gradio as gr
import os

import cohereAPI



# Model configurations
COHERE_MODELS = [
    "command-a-03-2025",
    "command-r7b-12-2024",
    "command-r-plus-08-2024",
    "command-r-08-2024",
    "command-light",
    "command-light-nightly",
    "command",
    "command-nightly"
]

def update_model_choices(provider):
    """Update model dropdown choices based on selected provider"""
    if provider == "Cohere":
        return gr.Dropdown(choices=COHERE_MODELS, value=COHERE_MODELS[0])
    else:
        return gr.Dropdown(choices=[], value=None)

def show_model_change_info(model_name):
    """Show info modal when model is changed"""
    if model_name:
        gr.Info(f"picking up from here with {model_name}")
    return model_name


async def respond(message, history, model_name="command-a-03-2025"):
    """Generate streaming response using Cohere API"""
    
    # Convert Gradio history format to API format
    conversation_history = []
    if history:
        for entry in history:
            if isinstance(entry, dict):
                # Clean dict format - only keep role and content
                if "role" in entry and "content" in entry:
                    conversation_history.append({
                        "role": entry["role"], 
                        "content": entry["content"]
                    })
            elif isinstance(entry, (list, tuple)) and len(entry) == 2:
                # Old format: [user_msg, assistant_msg]
                user_msg, assistant_msg = entry
                if user_msg:
                    conversation_history.append({"role": "user", "content": str(user_msg)})
                if assistant_msg:
                    conversation_history.append({"role": "assistant", "content": str(assistant_msg)})
            else:
                # Handle other formats gracefully
                continue
    
    # Get API key from environment
    api_key = os.getenv('COHERE_API_KEY')
    if not api_key:
        yield "Error: COHERE_API_KEY environment variable not set"
        return
    
    # System message for the chatbot
    system_message = """You are a helpful AI assistant. Provide concise but complete responses. 
                        Be direct and to the point while ensuring you fully address the user's question or request. 
                        Do not repeat the user's question in your response. Do not exceed 50 words."""

    try:
        # Use async streaming function
        partial_message = ""
        async for chunk in cohereAPI.send_message_stream_async(
            system_message=system_message,
            user_message=message,
            conversation_history=conversation_history,
            api_key=api_key,
            model_name=model_name
        ):
            partial_message += chunk
            yield partial_message
    except Exception as e:
        yield f"Error: {str(e)}"

with gr.Blocks() as demo:
    gr.Markdown("## Modular Chatbot")

    with gr.Row():
        with gr.Column(scale=2):
            chat_interface = gr.ChatInterface(
                fn=respond,
                type="messages",
                save_history=True
            )
            with gr.Accordion("Chat Settings", elem_id="chat_settings_group"):
                with gr.Row():
                    with gr.Column(scale=3):
                        provider = gr.Dropdown(
                            info="Provider",
                            choices=["Cohere", "OpenAI", "Anthropic", "Google", "HuggingFace"],
                            value="Cohere",
                            elem_id="provider_dropdown",
                            interactive=True,
                            show_label=False
                        )
                        model = gr.Dropdown(
                            info="Model",
                            choices=COHERE_MODELS,
                            value=COHERE_MODELS[0],
                            elem_id="model_dropdown",
                            interactive=True,
                            show_label=False
                        )
        
                    # Set up event handler for provider change
                    provider.change(
                        fn=update_model_choices,
                        inputs=[provider],
                        outputs=[model]
                    )
                    
                    # Set up event handler for model change
                    model.change(
                        fn=show_model_change_info,
                        inputs=[model],
                        outputs=[model]
                    )
                    
                    with gr.Column(scale=1):
                        temperature = gr.Slider(
                            label="Temperature",
                            info="Higher values make output more creative",
                            minimum=0.0,
                            maximum=1.0,
                            value=0.7,
                            step=0.01,
                            elem_id="temperature_slider",
                            interactive=True,
                            
                        )
                        max_tokens = gr.Textbox(
                            label="Max Tokens",
                            info="Higher values allow longer responses.",
                            value="8192",
                            elem_id="max_tokens_input",
                            interactive=True,
                            show_label=False,
                        )
        
if __name__ == "__main__":
    demo.launch()