import gradio as gr
import requests
from rembg import remove
from PIL import Image

# Available models
MODELS = [
    "openai-large", "phi", "qwen", "deepseek", "deepseek-reasoning",
    "hormoz", "llamascout", "gpt-3.5-turbo", "gpt-4", "llama-2-7b",
    "llama-2-13b", "llama-2-70b", "claude-2", "claude-instant", "palm-2",
    "code-llama", "mistral-7b", "mixtral-8x7b", "stablelm-tuned", "dolly-v2",
    "bloom", "gpt-neox"
]

# Custom CSS
css = """
@import url('https://fonts.googleapis.com/css2?family=Geist+Sans:wght@400;500&family=Syne:wght@600&display=swap');
.gradio-container {
    font-family: 'Geist Sans', sans-serif;
}
h1, .chatbot .label {
    font-family: 'Syne', sans-serif !important;
    font-weight: 600 !important;
}
.message p, .message .markdown {
    color: #000000 !important;  
}
.message.user {
    background: #f3f4f6 !important;
}
.message.bot {
    background: #ffffff !important;
    border: 1px solid #e5e7eb !important;
}
"""

def get_ai_response(prompt, model):
    try:
        url = f"https://text.pollinations.ai/{prompt}?model={model}"
        response = requests.get(url, timeout=30)
        return response.text
    except Exception as e:
        return f"⚠ Error: {str(e)}"

def chat(message, history, model, image=None):
    if message.strip() == "/bgremover" and image is not None:
        try:
            # Remove background using rembg 
            output_image = remove(image.convert("RGBA"))
            history.append((message, (output_image,)))
        except Exception as e:
            history.append((message, f"❌ Background removal failed: {str(e)}"))
    else:
        bot_response = get_ai_response(message, model)
        history.append((message, bot_response))
    
    return "", history, None  # Clear image input after processing

with gr.Blocks(css=css, title="✦ AI Chat") as demo:
    gr.Markdown("# ✦ AI Chat")
    
    with gr.Row():
        model_dropdown = gr.Dropdown(
            choices=MODELS,
            value=MODELS[0],
            label="Select Model",
            interactive=True
        )
    
    chatbot = gr.Chatbot(
        avatar_images=(None, "https://pomf2.lain.la/f/8u6md29u.png"), 
        bubble_full_width=False,
        height=500
    )
    
    with gr.Row():
        msg = gr.Textbox(
            placeholder="Type your message here...",
            container=False,
            scale=3,
            lines=3
        )
        image_input = gr.Image(type="pil", label="Upload Image")
        submit = gr.Button("➜", variant="primary", scale=1)
    
    with gr.Row():
        clear_btn = gr.Button("🗑️ Clear")
        restart_btn = gr.Button("🔄 Restart")

    # Event handlers
    msg.submit(chat, [msg, chatbot, model_dropdown, image_input], [msg, chatbot, image_input])
    submit.click(chat, [msg, chatbot, model_dropdown, image_input], [msg, chatbot, image_input])
    clear_btn.click(lambda: None, None, chatbot, queue=False)
    restart_btn.click(lambda: None, None, chatbot, queue=False)

demo.launch()