import gradio as gr
from huggingface_hub import InferenceClient

import os

CHAT_MODEL = "Qwen/Qwen2.5-7B-Instruct-1M:featherless-ai"

HF_API_KEY = os.getenv("HF_API_KEY")

# ── Tab 1: Get weather  ────────────────────────────────────────────────────────────
def get_weather(location: str) -> str:
    """
    Retrieves the current weather for a given location.
    Args:
        location: The city or place for which to fetch the weather.
    """
    # In a real tool, call a weather API here.
    return f"The weather in {location} is sunny with mild temperatures."

tab1 = gr.Interface(
    fn=get_weather,
    inputs=[
        gr.Textbox(label="Location", placeholder="e.g. Brussels"),
    ],
    outputs=gr.Textbox(label="Weather"),
    title="Get Weather",
    examples=[["Brussels"], ["London"]],
)

# ── Tab 2: Ask an LLM ─────────────────────────────────────────────────────────
llm_client = InferenceClient(model=CHAT_MODEL, api_key=HF_API_KEY)

def ask_llm(prompt, system_prompt, max_tokens):
    response = llm_client.chat.completions.create(
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt},
        ],
        max_tokens=int(max_tokens),
    )
    return response.choices[0].message.content

tab2 = gr.Interface(
    fn=ask_llm,
    inputs=[
        gr.Textbox(label="Your question", lines=3),
        gr.Textbox(label="System prompt", value="You are a helpful assistant.", lines=2),
        gr.Slider(50, 500, value=200, step=50, label="Max tokens"),
    ],
    outputs=gr.Textbox(label="Answer", lines=8),
    title="Ask an LLM",
)

# ── Tab 3: Image Generator ────────────────────────────────────────────────────
img_client = InferenceClient("black-forest-labs/FLUX.1-schnell", api_key=HF_API_KEY)

def generate_image(prompt, style):
    """
    Generate an image from a text description using Imagen via Google GenAI.
    Parameters
    ----------
    prompt : str
        A natural-language description of the image to generate.
    Returns
    -------
    PIL.Image.Image
        The generated image as a Pillow Image object, ready to display in Gradio.
    """

    full_prompt = f"{prompt}, {style} style" if style else prompt
    return img_client.text_to_image(full_prompt)

tab3 = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Describe your image", placeholder="A futuristic city at sunset...", lines=2),
        gr.Radio(
            ["photorealistic", "watercolor painting", "pixel art", "pencil sketch"],
            label="Style",
            value="photorealistic",
        ),
    ],
    outputs=gr.Image(label="Generated image"),
    title="Image Generator",
)

# ── Tab 4: Multi-turn Chat ────────────────────────────────────────────────────
chat_client = InferenceClient(model=CHAT_MODEL, api_key=HF_API_KEY)

def chat_fn(message, history):
    messages = history + [{"role": "user", "content": message}]
    response = chat_client.chat.completions.create(messages=messages)
    return response.choices[0].message.content

tab4 = gr.ChatInterface(
    fn=chat_fn,
    title="Chat with Qwen",
    examples=["Tell me a joke", "Explain AI in simple terms", "What is Gradio?"],
)

# ── Assemble ──────────────────────────────────────────────────────────────────
app = gr.TabbedInterface(
    [tab1, tab2, tab3, tab4],
    tab_names=["Get Weather", "Ask LLM", "Image Gen", "Chat"],
)
app.launch(mcp_server=True)