import gradio as gr
import cohere
import os
import uuid

cohere_api_key = os.getenv("COHERE_API_KEY")
co = cohere.Client(cohere_api_key, client_name="huggingface-rp")

def trigger_example(example: str) -> tuple:
    chat, updated_history = generate_response(example)
    return chat, updated_history

def generate_response(user_message: str, cid: str, history: list = None) -> tuple:
    if history is None:
        history = []
    if not cid:  # Simplified cid check
        cid = str(uuid.uuid4())
    history.append(user_message)
    stream = co.chat_stream(
        message=user_message, conversation_id=cid, model='command-r-plus', 
        connectors=[{"id": "web-search"}], temperature=0.3
    )
    output = ""
    for idx, response in enumerate(stream):
        if response.event_type == "text-generation":
            output += response.text
        if idx == 0:
            history.append(" " + output)
        else:
            history[-1] = output
        chat = [(history[i].strip(), history[i + 1].strip()) for i in range(0, len(history) - 1, 2)]
        yield chat, history, cid
    return chat, history, cid

def clear_chat() -> tuple:
    return [], [], str(uuid.uuid4())

custom_css = """
.gradio-container {
    max-width: 800px;
    margin: 40px auto;
    padding: 20px;
    border: 1px solid #ddd;
    border-radius: 10px;
    box-shadow: 0 0 10px rgba(0, 0, 0, 0.1);
}

.gradio-input {
    margin-bottom: 20px;
}

.gradio-button {
    background-color: #4CAF50;
    color: #fff;
    padding: 10px 20px;
    border: none;
    border-radius: 5px;
    cursor: pointer;
}

.gradio-button:hover {
    background-color: #3e8e41;
}

#chat-message {
    font-size: 14px;
    min-height: 300px;
    padding: 10px;
    border: 1px solid #ddd;
    border-radius: 10px;
}
"""

with gr.Blocks(analytics_enabled=False, css=custom_css) as demo:
    cid = gr.State("")
    history = gr.State([])

    with gr.Row():
        with gr.Column(scale=2):
            gr.Markdown("""
            **Multi lingual law Assistant**
            ===============
            Multi lingual Law assistant for Indian legal terms and corpus
            """)
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(show_label=False, show_share_button=False, show_copy_button=True)

    with gr.Row():
        user_message = gr.Textbox(lines=1, placeholder="Ask anything...", label="Input", show_label=False)
        submit_button = gr.Button("Submit")
        clear_button = gr.Button("Clear chat")

    user_message.submit(fn=generate_response, inputs=[user_message, cid, history], outputs=[chatbot, history, cid], concurrency_limit=32)
    submit_button.click(fn=generate_response, inputs=[user_message, cid, history], outputs=[chatbot, history, cid], concurrency_limit=32)
    clear_button.click(fn=clear_chat, inputs=None, outputs=[chatbot, history, cid], concurrency_limit=32)

if __name__ == "__main__":
    demo.queue(api_open=False, max_size=40).launch(show_api=False)