import gradio as gr
from huggingface_hub import InferenceClient
import time

# Clientes para texto e imagem
chat_client = InferenceClient("lambdaindie/lambdai")
image_client = InferenceClient("stabilityai/stable-diffusion-2")

# Fonte global
gr.themes.Base().set(font=["JetBrains Mono", "monospace"])

css = """
body {
    font-family: 'JetBrains Mono', monospace;
    background-color: #111;
    color: #e0e0e0;
}
.gr-textbox textarea {
    background-color: #181818 !important;
    color: #fff !important;
    font-family: 'JetBrains Mono', monospace;
    border-radius: 8px;
}
.markdown-think {
    background-color: #1e1e1e;
    border-left: 4px solid #555;
    padding: 10px;
    margin-bottom: 8px;
    font-style: italic;
    animation: pulse 1.5s infinite ease-in-out;
}
@keyframes pulse {
    0% { opacity: 0.6; }
    50% { opacity: 1.0; }
    100% { opacity: 0.6; }
}
"""

# Função do chatbot com raciocínio
def respond(message, history, system_message, max_tokens, temperature, top_p):
    messages = [{"role": "system", "content": system_message}] if system_message else []

    for user, assistant in history:
        if user:
            messages.append({"role": "user", "content": user})
        if assistant:
            messages.append({"role": "assistant", "content": assistant})

    thinking_prompt = messages + [
        {"role": "user", "content": f"{message}\n\nThink step-by-step before answering."}
    ]

    reasoning = ""
    yield '<div class="markdown-think">Thinking...</div>'

    for chunk in chat_client.chat_completion(
        thinking_prompt,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = chunk.choices[0].delta.content or ""
        reasoning += token
        yield f'<div class="markdown-think">{reasoning.strip()}</div>'

    time.sleep(0.5)

    final_prompt = messages + [
        {"role": "user", "content": message},
        {"role": "assistant", "content": reasoning.strip()},
        {"role": "user", "content": "Now answer based on your reasoning above."}
    ]

    final_answer = ""
    for chunk in chat_client.chat_completion(
        final_prompt,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = chunk.choices[0].delta.content or ""
        final_answer += token
        yield final_answer.strip()

# Função para gerar imagem
def generate_image(prompt):
    return image_client.text_to_image(prompt, guidance_scale=7.5)

# Interface Gradio
with gr.Blocks(css=css, theme=gr.themes.Base()) as demo:
    gr.Markdown("# λmabdAI")

    with gr.Tabs():
        with gr.Tab("Chat"):
            gr.ChatInterface(
                fn=respond,
                additional_inputs=[
                    gr.Textbox(
                        value="You are a concise, logical AI that explains its reasoning clearly before answering.",
                        label="System Message"
                    ),
                    gr.Slider(64, 2048, value=512, step=1, label="Max Tokens"),
                    gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature"),
                    gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")
                ]
            )

        with gr.Tab("Image Generator"):
            gr.Markdown("### Generate an image from a prompt")
            prompt = gr.Textbox(label="Prompt")
            output = gr.Image(type="pil")
            btn = gr.Button("Generate")
            btn.click(fn=generate_image, inputs=prompt, outputs=output)

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