import time

import gradio as gr
import aiohttp
import asyncio
import json

# Define the types and their associated models with labels and values
types_and_models = {
    'groq': [
        ('Llama 3.1 8B Instant', 'llama-3.1-8b-instant'),
        ('Llama 3.1 70B Versatile', 'llama-3.1-70b-versatile'),
        ('Llama 3.2 3B Instant', 'llama-3.2-3b-preview'),
    ],
    'openai': [
        ('GPT-4o Mini', 'gpt-4o-mini'),
        ('GPT-4o', 'gpt-4o'),
    ],
    'mistral': [
        ('Mistral Large', 'mistral-large-latest'),
        ('Mistral small', 'mistral-small-latest'),
    ],
    'deepseek': [
        ('DeepSeek V2.5', 'deepseek-chat'),
    ],
    'deepinfra': [
        ('Llama 8B', 'meta-llama/Meta-Llama-3.1-8B-Instruct'),
        ('Llama 3B', 'meta-llama/Llama-3.2-3B-Instruct'),
        ('Llama 70B', 'meta-llama/Meta-Llama-3.1-70B-Instruct'),
        ('Llama 405B', 'meta-llama/Meta-Llama-3.1-405B-Instruct'),
        ('Qwen 72B', 'Qwen/Qwen2.5-72B-Instruct'),
    ],

    'custom': [],
    'open_router': [],
}


def update_model_input(selected_type):
    if selected_type in ['custom', 'open_router']:
        return gr.update(visible=False), gr.update(visible=True)
    else:
        choices = types_and_models[selected_type]
        default_value = choices[0][1] if choices else None
        return gr.update(choices=choices, value=default_value, visible=True), gr.update(visible=False)


async def chat_with_server(chat_history, api_key, base_url, selected_type, selected_model, custom_model, temperature,
                           max_tokens):
    headers = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {api_key}',
    }

    api_url = 'https://api.logicboost.net'  # API URL

    conversation = []
    for user_msg, assistant_msg in reversed(chat_history):
        if user_msg is not None:
            conversation.append(f"User: {user_msg}")
        if assistant_msg and not assistant_msg.startswith("Assistant is thinking"):
            conversation.append(f"Assistant: {assistant_msg}")
        elif assistant_msg and assistant_msg.startswith("Assistant is thinking"):
            conversation.append(f"System: {assistant_msg}")

    user_content = "\n".join(conversation)
    api_messages = [{'role': 'user', 'content': user_content}]

    model = custom_model if selected_type in ['custom', 'open_router'] else selected_model

    payload = {
        'type': selected_type,
        'baseUrl': base_url,
        'model': model,
        'messages': api_messages,
        'temperature': temperature,
        'max_tokens': int(max_tokens),
        'stream': True,
    }

    start_time = time.time()

    if chat_history:
        chat_history[-1][1] = "Assistant is thinking (0 sec)"
        yield chat_history

    async with aiohttp.ClientSession() as session:
        try:
            response_task = asyncio.create_task(session.post(
                f'{api_url}/v1/chat/completions',
                headers=headers,
                json=payload,
                timeout=1200
            ))

            dots_count = 0
            assistant_message = ''

            while not response_task.done():
                elapsed_time = time.time() - start_time
                minutes, seconds = divmod(int(elapsed_time), 60)
                time_str = f"{minutes} min {seconds} sec" if minutes > 0 else f"{seconds} sec"

                dots = '.' * ((dots_count % 3) + 1)
                dots_count += 1
                if chat_history:
                    chat_history[-1][1] = f"Assistant is thinking ({time_str}){dots}"
                    yield chat_history
                await asyncio.sleep(0.5)

            response = await response_task
            if response.status != 200:
                error_message = f"Error {response.status}: {await response.text()}"
                chat_history[-1][1] = error_message
                yield chat_history
                return

            assistant_message = ''
            async for line in response.content:
                line = line.decode('utf-8').strip()
                if line == '':
                    continue
                if line == 'data: [DONE]':
                    break
                if line.startswith('data: '):
                    data_str = line[6:].strip()
                    if data_str:
                        data = json.loads(data_str)
                        if 'choices' in data and len(data['choices']) > 0:
                            delta = data['choices'][0]['delta']
                            if 'content' in delta:
                                content = delta['content']
                                assistant_message += content
                                chat_history[-1][1] = assistant_message
                                yield chat_history
            yield chat_history

        except Exception as e:
            error_message = f"Error: {str(e)}"
            chat_history[-1][1] = error_message
            yield chat_history


def reset_chat():
    return []


with gr.Blocks() as demo:
    gr.Markdown("""# 🧠 XoT Self-Dialog
    It's an alternative realization of the analog - GPT-o1""")

    with gr.Row():
        api_key = gr.Textbox(label="API Key", placeholder="Enter your API Key",
                             value="ggkjgf", type="password", lines=1)
        base_url = gr.Textbox(label="Base URL", placeholder="https://api.groq.com/openai/v1", lines=2)
        selected_type = gr.Dropdown(
            label="Provider",
            choices=list(types_and_models.keys()),
            value=list(types_and_models.keys())[0]
        )
        # Initialize the selected_model with choices and value
        initial_type = list(types_and_models.keys())[0]
        initial_choices = types_and_models[initial_type]
        selected_model = gr.Dropdown(
            label="Model",
            choices=initial_choices,
            value=initial_choices[0][1] if initial_choices else None
        )
        custom_model = gr.Textbox(label="Custom Model", placeholder="Enter custom model name",
                                  visible=False)
        temperature = gr.Slider(label="Temperature", minimum=0, maximum=1, value=0.8, step=0.1)
        max_tokens = gr.Number(label="Max Tokens", value=4096)

    # Initialize the chatbot component
    chatbot = gr.Chatbot(elem_id="chatbot", height=400, render_markdown=True, latex_delimiters=[
        {"left": "$", "right": "$", "display": False}, {"left": "\\[", "right": "\\]", "display": True}])

    with gr.Row():
        with gr.Column(scale=10):
            user_input = gr.Textbox(
                show_label=False,
                placeholder="Type your message here...",
                lines=5,
                container=False
            )
        with gr.Column(min_width=50, scale=1):
            send_button = gr.Button("Send", variant="primary")
            reset_button = gr.Button("Reset Chat")

    # Update the models dropdown when the type changes
    selected_type.change(
        update_model_input,
        inputs=selected_type,
        outputs=[selected_model, custom_model]
    )


    def user(user_message, chat_history):
        # Append the user's message to the chat history
        if not chat_history:
            chat_history = []
        chat_history = chat_history + [[user_message, None]]
        return user_message, chat_history


    # Set up the event handlers
    send_button.click(
        user,
        inputs=[user_input, chatbot],
        outputs=[user_input, chatbot]
    ).then(
        chat_with_server,
        inputs=[chatbot, api_key, base_url, selected_type, selected_model, custom_model, temperature, max_tokens],
        outputs=chatbot,
    )

    reset_button.click(reset_chat, outputs=chatbot)

    # Include CSS to format code blocks and other markdown elements
    gr.HTML("""
    <style>
    .message pre {
        background-color: #f0f0f0;
        padding: 10px;
        border-radius: 5px;
        overflow-x: auto;
    }
    .message code {
        background-color: #f0f0f0;
        padding: 2px 4px;
        border-radius: 3px;
    }
    
    .description-box {
        background-color: #f0f8ff;  /* Light blue background */
        border: 1px solid #add8e6;  /* Light blue border */
        border-radius: 10px;
        padding: 20px;
        margin-top: 20px;
    }
    </style>
    """)

    gr.Markdown("""
    <div class="description-box">

    ## 🧠 About XoT Self-Dialog

    XoT Self-Dialog is a modification of the **XoT** (_Everything of Thoughts_) technique, designed to enhance LLM reasoning. This implementation incorporates self-dialog and automated Monte Carlo Tree Search (MCTS) to improve problem-solving capabilities.

    The technique aims to refine the thought process of LLMs, allowing for more effective exploration of solution spaces and the generation of multiple solution paths. By leveraging self-dialog and MCTS, XoT Self-Dialog enables LLMs to tackle complex, multi-solution problems more efficiently.

    This interface demonstrates the application of XoT Self-Dialog, showcasing its potential to enhance reasoning across various problem domains.

    </div>
    """)

demo.queue()
demo.launch()