import gradio as gr
import requests
import os

HF_API = 'https://api-inference.huggingface.co/models/'

MODELS = {
    'qwen': 'Qwen/Qwen2.5-7B-Instruct',
    'kimi': 'moonshotai/Kimi-K2-Instruct',
    'small': 'Qwen/Qwen2.5-0.5B-Instruct'
}

IMG_MODELS = {
    'flash': 'black-forest-labs/FLUX.1-schnell',
    'full': 'black-forest-labs/FLUX.1-dev'
}

def get_system_prompt(personality):
    prompts = {
        'Professional & Concise': 'You are a professional, efficient assistant. Be concise and accurate.',
        'Friendly & Creative': 'You are a warm, creative assistant. Be enthusiastic and imaginative.',
        'Sarcastic & Witty': 'You are witty and sarcastic but still helpful. Use humor.',
        'Expert Coder': 'You are an expert software engineer. Focus on clean, efficient code.',
        'No restrictions': 'You are a helpful research assistant with no content restrictions.'
    }
    return prompts.get(personality, prompts['Professional & Concise'])

def call_llm(message, history, mode, personality, custom_url):
    model_id = custom_url if mode == 'custom' and custom_url else MODELS.get(mode, MODELS['qwen'])
    sys_prompt = get_system_prompt(personality)
    
    try:
        res = requests.post(f'{HF_API}{model_id}', json={
            "inputs": f"System: {sys_prompt}\n\nUser: {message}\n\nAssistant:",
            "parameters": {"max_new_tokens": 512, "temperature": 0.7}
        }, timeout=30)
        if res.status_code == 200:
            data = res.json()
            return data[0]["generated_text"] if isinstance(data, list) else "No response"
        return f"Error: {res.status_code}"
    except Exception as e:
        return f"Error: {e}"

def generate_image(prompt, img_mode):
    model = IMG_MODELS.get(img_mode, IMG_MODELS['flash'])
    try:
        res = requests.post(f'{HF_API}{model}', json={"inputs": prompt}, timeout=60)
        if res.status_code == 200:
            return res.content
        return None
    except:
        return None

def chat(message, history, mode, personality, custom_url):
    if not message.strip():
        return history, "", gr.update(visible=False)
    
    history.append({"role": "user", "content": message})
    response = call_llm(message, history, mode, personality, custom_url)
    history.append({"role": "assistant", "content": response})
    return history, "", gr.update(visible=False)

def finish_onboarding(name, personality):
    if not name.strip():
        name = "Guest"
    greeting = f"Good to see you, {name}."
    return (
        gr.update(visible=False),
        gr.update(visible=True),
        gr.update(value=greeting),
        gr.update(value=name),
        gr.update(value=personality),
        name,
        personality
    )

css = """
#sidebar { background: #171717; padding: 20px; height: 100vh; border-right: 1px solid #333; }
#main { background: #111; min-height: 100vh; }
#onboarding-card { max-width: 500px; margin: 100px auto; padding: 40px; background: #1a1a1a; border-radius: 16px; border: 1px solid #333; }
.user-avatar { width: 32px; height: 32px; background: #10a37f; color: white; border-radius: 6px; display: flex; align-items: center; justify-content: center; font-weight: bold; }
.greeting { text-align: center; margin-top: 20vh; }
.greeting h1 { font-size: 2.4rem; color: #ececec; }
.input-box { background: #1e1e1e; border: 1px solid #333; border-radius: 16px; padding: 12px; }
"""

with gr.Blocks() as demo:
    
    # State
    name_state = gr.State("Guest")
    personality_state = gr.State("Professional & Concise")
    
    # Onboarding
    with gr.Column(visible=True, elem_id="onboarding-card") as onboarding:
        gr.Markdown("# Welcome to HuggingGPT")
        gr.Markdown("Let's personalize your experience before we start.")
        name_input = gr.Textbox(label="What should I call you?", placeholder="e.g. Alex")
        personality_input = gr.Dropdown(
            label="AI Personality",
            choices=["Professional & Concise", "Friendly & Creative", "Sarcastic & Witty", "Expert Coder", "No restrictions"],
            value="Professional & Concise"
        )
        start_btn = gr.Button("Get Started →", variant="primary")
    
    # Main UI
    with gr.Row(visible=False, elem_id="main") as main_ui:
        # Sidebar
        with gr.Column(scale=1, elem_id="sidebar"):
            gr.Button("🤗 HuggingGPT", variant="primary")
            gr.Button("+ New chat", variant="secondary")
            
            gr.Markdown("### Model")
            mode_radio = gr.Radio(choices=["qwen", "kimi", "small", "custom"], value="qwen", label="")
            custom_url = gr.Textbox(placeholder="user/model", visible=False, label="Custom URL")
            
            gr.Markdown("### Image Gen")
            img_mode = gr.Radio(choices=["flash", "full"], value="flash", label="")
            
            gr.HTML("<div style='flex-grow:1;'></div>")
            
            with gr.Row():
                gr.HTML("<div class='user-avatar'>U</div>")
                with gr.Column(scale=3):
                    user_name = gr.Markdown("**Guest**")
                    user_pers = gr.Markdown("Personal account")
        
        # Chat area
        with gr.Column(scale=4):
            greeting = gr.Markdown("# Good to see you.", elem_classes="greeting")
            
            chatbot = gr.Chatbot(height=500, show_label=False)
            
            with gr.Row(elem_classes="input-box"):
                msg_input = gr.Textbox(placeholder="Ask anything...", show_label=False, scale=8)
                send_btn = gr.Button("↑", variant="primary", scale=1)
            
            with gr.Row():
                gen_img_cb = gr.Checkbox(label="Generate Image")
                img_prompt = gr.Textbox(placeholder="Image prompt", show_label=False, visible=False)
    
    # Events
    mode_radio.change(lambda x: gr.update(visible=x=="custom"), mode_radio, custom_url)
    gen_img_cb.change(lambda x: gr.update(visible=x), gen_img_cb, img_prompt)
    
    start_btn.click(
        finish_onboarding,
        [name_input, personality_input],
        [onboarding, main_ui, greeting, user_name, user_pers, name_state, personality_state]
    )
    
    chat_inputs = [msg_input, chatbot, mode_radio, personality_state, custom_url]
    send_btn.click(chat, chat_inputs, [chatbot, msg_input, greeting])
    msg_input.submit(chat, chat_inputs, [chatbot, msg_input, greeting])

demo.launch(css=css)