# ==========================================
# 🖍️ SHINCHAN AI — Gradio + Ollama
# Compatible with Gradio 5.x AND 6.x
# ==========================================

import gradio as gr
import requests
import json

OLLAMA_URL = "http://localhost:11434"
MODEL_NAME = "shinchan-ai"


def check_ollama_health():
    """Check if Ollama server is running and model is available."""
    try:
        r = requests.get(f"{OLLAMA_URL}/api/tags", timeout=5)
        if r.status_code == 200:
            models = [m["name"] for m in r.json().get("models", [])]
            return MODEL_NAME in " ".join(models)
    except Exception:
        pass
    return False


def user_submit(message, history):
    """Add user message to chat history."""
    if not message or not message.strip():
        return "", history, history
    history = history + [{"role": "user", "content": message.strip()}]
    return "", history, history


def bot_respond(history):
    """Stream Shinchan's response from Ollama."""
    if not history:
        return history, history

    if not check_ollama_health():
        history = history + [{
            "role": "assistant",
            "content": "😤 Arey yaar! Mera dimaag abhi load ho raha hai... 2 minute ruko aur phir try karo! 😴"
        }]
        yield history, history
        return

    messages = []
    for h in history:
        messages.append({"role": h["role"], "content": h["content"]})

    payload = {
        "model": MODEL_NAME,
        "messages": messages,
        "stream": True,
        "options": {
            "temperature": 0.8,
            "top_p": 0.9,
            "num_predict": 300,
        }
    }

    history = history + [{"role": "assistant", "content": ""}]

    try:
        response = requests.post(
            f"{OLLAMA_URL}/api/chat",
            json=payload,
            stream=True,
            timeout=300,
        )
        response.raise_for_status()

        for line in response.iter_lines():
            if line:
                try:
                    chunk = json.loads(line)
                    if "message" in chunk:
                        token = chunk["message"].get("content", "")
                        if token:
                            history[-1]["content"] += token
                            yield history, history
                except json.JSONDecodeError:
                    continue

        if not history[-1]["content"].strip():
            history[-1]["content"] = (
                "Arey yaar, dimag mein kuch aaya hi nahi! 😤 "
                "Phir se pucho na... 😏"
            )
            yield history, history

    except requests.exceptions.Timeout:
        history[-1]["content"] = (
            "😴 Bohut time lag gaya yaar... Mummy ne internet slow kar diya! "
            "Phir se try karo! 😤"
        )
        yield history, history

    except Exception as e:
        history[-1]["content"] = f"Oops! Error aa gaya 😤: {str(e)}"
        yield history, history


def clear_chat():
    return [], []


# ==========================================
# 🎨 BUILD GRADIO UI (Compatible with v5 & v6)
# ==========================================

# Detect Gradio version
GRADIO_VERSION = int(gr.__version__.split(".")[0])
print(f"📦 Detected Gradio version: {gr.__version__} (major={GRADIO_VERSION})")

# Build kwargs based on version
blocks_kwargs = {}
chatbot_kwargs = {
    "height": 460,
    "label": "Shinchan 🖍️",
    "show_copy_button": True,
}
launch_kwargs = {
    "server_name": "0.0.0.0",
    "server_port": 7860,
    "share": False,
}

if GRADIO_VERSION >= 6:
    # Gradio 6.x: theme/css go in launch(), type removed from Chatbot
    launch_kwargs["theme"] = gr.themes.Soft()
    launch_kwargs["css"] = """
        footer { display: none !important; }
        .gr-button { border-radius: 12px !important; }
    """
else:
    # Gradio 5.x and below: theme/css in Blocks(), type in Chatbot
    blocks_kwargs["theme"] = gr.themes.Soft()
    blocks_kwargs["css"] = """
        footer { display: none !important; }
        .gr-button { border-radius: 12px !important; }
    """
    chatbot_kwargs["type"] = "messages"


with gr.Blocks(title="Shinchan AI 🖍️", **blocks_kwargs) as demo:

    # --- Header ---
    gr.HTML("""
    <div style="text-align:center; padding:20px 10px 5px;">
        <h1 style="margin-bottom:4px;">🖍️ Shinchan AI Chat</h1>
        <p style="font-size:16px; color:#555;">
            Shin-chan se baat karo! Chocobi nahi milegi 😏😂
        </p>
        <p style="font-size:11px; color:#999;">
            Powered by ShinchanAI-small · Ollama · 🤗 Hugging Face Spaces
        </p>
    </div>
    """)

    # --- Chatbot ---
    chatbot = gr.Chatbot(**chatbot_kwargs)

    # --- Input ---
    with gr.Row():
        msg = gr.Textbox(
            placeholder="Shinchan se kuch pucho... 😋",
            label="Your Message",
            container=False,
            scale=7,
        )
        submit_btn = gr.Button("Send 🚀", variant="primary", scale=1)

    with gr.Row():
        clear_btn = gr.Button("🗑️ Clear Chat")

    gr.HTML("""
    <p style="text-align:center; color:gray; margin-top:8px; font-size:13px;">
        Try:
        <b>Tu kaun hai?</b> &nbsp;|&nbsp;
        <b>Chocobi de do</b> &nbsp;|&nbsp;
        <b>School jaana hai?</b> &nbsp;|&nbsp;
        <b>Mujhe coding sikha do</b> &nbsp;|&nbsp;
        <b>Action Kamen kaisa hai?</b> &nbsp;|&nbsp;
        <b>Nanako miss kahan hai?</b>
    </p>
    """)

    # --- State ---
    chat_history = gr.State([])

    # --- Wire Events ---
    msg.submit(
        user_submit,
        inputs=[msg, chat_history],
        outputs=[msg, chatbot, chat_history],
    ).then(
        bot_respond,
        inputs=[chat_history],
        outputs=[chatbot, chat_history],
    )

    submit_btn.click(
        user_submit,
        inputs=[msg, chat_history],
        outputs=[msg, chatbot, chat_history],
    ).then(
        bot_respond,
        inputs=[chat_history],
        outputs=[chatbot, chat_history],
    )

    clear_btn.click(clear_chat, outputs=[chatbot, chat_history])


# ==========================================
# 🚀 LAUNCH
# ==========================================
print("🎨 Gradio UI launching on port 7860...")

demo.queue()
demo.launch(**launch_kwargs)