import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
from threading import Thread
import spaces
import json
import time
from datetime import datetime

# ── Model ──────────────────────────────────────────────────────────────────────
MODEL_ID = "huihui-ai/Huihui-Qwen3.5-2B-abliterated"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
model.eval()

# ── Presets ────────────────────────────────────────────────────────────────────
SYSTEM_PRESETS = {
    "🤖 Varsayılan Asistan": "You are a helpful, harmless and honest assistant.",
    "💻 Kod Uzmanı": "You are an expert software engineer. Write clean, efficient, well-documented code. Always explain your reasoning.",
    "✍️ Yaratıcı Yazar": "You are a creative writing assistant with a vivid imagination. Help craft compelling stories, characters, and narratives.",
    "🔬 Bilim Danışmanı": "You are a knowledgeable science advisor. Explain complex topics clearly with accurate information and real-world examples.",
    "🗣️ Türkçe Asistan": "Sen yardımsever, bilgili bir Türkçe asistansın. Her zaman Türkçe yanıt ver ve net, anlaşılır açıklamalar yap.",
    "🎯 Özel": "",
}

# ── GPU generation (streaming) ────────────────────────────────────────────────
@spaces.GPU
def generate_stream(message, history, system_prompt, max_new_tokens, temperature, top_p, repetition_penalty):
    messages = []
    if system_prompt.strip():
        messages.append({"role": "system", "content": system_prompt})
    for h in history:
        # Yeni Gradio: ChatMessage dict {role, content} | Eski: [user, bot]
        if isinstance(h, dict):
            if h.get("content"):
                messages.append({"role": h["role"], "content": h["content"]})
        else:
            if h[0]:
                messages.append({"role": "user",      "content": str(h[0])})
            if h[1]:
                messages.append({"role": "assistant", "content": str(h[1])})
    messages.append({"role": "user", "content": message})

    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer([text], return_tensors="pt").to(model.device)

    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    gen_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        do_sample=temperature > 0.01,
    )

    thread = Thread(target=model.generate, kwargs=gen_kwargs)
    thread.start()

    partial = ""
    for chunk in streamer:
        partial += chunk
        yield partial

    thread.join()


# ── Helpers ────────────────────────────────────────────────────────────────────
def update_system_prompt(preset_name):
    return SYSTEM_PRESETS.get(preset_name, "")

def export_chat(history, system_prompt):
    if not history:
        return None
    data = {
        "exported_at": datetime.now().isoformat(),
        "model": MODEL_ID,
        "system_prompt": system_prompt,
        "messages": [
            {"user": h.get("content") if isinstance(h, dict) and h.get("role") == "user" else (h[0] if not isinstance(h, dict) else ""),
             "assistant": ""} if (isinstance(h, dict) and h.get("role") == "user") else
            {"assistant": h.get("content") if isinstance(h, dict) else h[1]}
            for h in history
        ],
    }
    path = f"/tmp/chat_{int(time.time())}.json"
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
    return path

def count_tokens(text):
    if not text:
        return 0
    return len(tokenizer.encode(text))

def get_stats(history):
    if not history:
        return '<div id="stats-bar">Henüz mesaj yok.</div>'
    total_msgs = len(history) * 2
    total_chars = sum(
        len(h.get("content", "") if isinstance(h, dict) else (h[0] or "") + (h[1] or ""))
        for h in history
    )
    return f'<div id="stats-bar">💬 {total_msgs} mesaj &nbsp;•&nbsp; 📝 {total_chars:,} karakter</div>'

def reset_parameters():
    return 1024, 0.7, 0.9, 1.1

def update_token_count(text):
    n = count_tokens(text)
    color = "#4ade80" if n < 512 else "#fbbf24" if n < 1024 else "#f87171"
    return f'<div id="token-info" style="color:{color}">Tokens: {n:,}</div>'

# ── CSS ────────────────────────────────────────────────────────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Syne:wght@400;500;600;700;800&family=JetBrains+Mono:wght@400;500&family=Inter:wght@300;400;500&display=swap');

:root {
    --bg-primary:    #0a0a0f;
    --bg-secondary:  #111118;
    --bg-tertiary:   #18181f;
    --bg-card:       #1c1c25;
    --border:        #2a2a38;
    --border-light:  #35354a;
    --accent:        #7c6af7;
    --accent-bright: #9d8fff;
    --accent-glow:   rgba(124, 106, 247, 0.15);
    --text-primary:  #e8e8f0;
    --text-secondary:#9090a8;
    --text-muted:    #55556a;
    --user-bg:       #1e1a3a;
    --bot-bg:        #141420;
    --success:       #4ade80;
    --warning:       #fbbf24;
    --danger:        #f87171;
    --radius:        12px;
    --radius-lg:     18px;
}

*, *::before, *::after { box-sizing: border-box; }

body, .gradio-container {
    font-family: 'Inter', sans-serif !important;
    background: var(--bg-primary) !important;
    color: var(--text-primary) !important;
    min-height: 100vh;
}

.gradio-container::before {
    content: '';
    position: fixed;
    top: -50%; left: -50%;
    width: 200%; height: 200%;
    background: radial-gradient(ellipse at 20% 20%, rgba(124,106,247,0.06) 0%, transparent 50%),
                radial-gradient(ellipse at 80% 80%, rgba(99,179,237,0.04) 0%, transparent 50%);
    pointer-events: none;
    z-index: 0;
    animation: bgShift 20s ease-in-out infinite alternate;
}
@keyframes bgShift {
    from { transform: translate(0,0) rotate(0deg); }
    to   { transform: translate(2%,2%) rotate(3deg); }
}

#header-block {
    background: linear-gradient(135deg, var(--bg-secondary) 0%, var(--bg-tertiary) 100%);
    border: 1px solid var(--border);
    border-radius: var(--radius-lg);
    padding: 28px 36px;
    margin-bottom: 20px;
    position: relative;
    overflow: hidden;
}
#header-block::before {
    content: '';
    position: absolute;
    top: 0; left: 0; right: 0; height: 2px;
    background: linear-gradient(90deg, transparent, var(--accent), var(--accent-bright), transparent);
}
#header-block h1 {
    font-family: 'Syne', sans-serif !important;
    font-size: 1.9rem !important;
    font-weight: 800 !important;
    letter-spacing: -0.03em !important;
    background: linear-gradient(135deg, #fff 0%, var(--accent-bright) 60%, #63b3ed 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    margin: 0 0 6px 0 !important;
}
#header-block p {
    color: var(--text-secondary) !important;
    font-size: 0.88rem !important;
    margin: 0 !important;
    font-weight: 300;
    letter-spacing: 0.01em;
}

#chatbot {
    background: var(--bg-secondary) !important;
    border: 1px solid var(--border) !important;
    border-radius: var(--radius-lg) !important;
    font-family: 'Inter', sans-serif !important;
    font-size: 0.92rem !important;
}
#chatbot .message.user {
    background: var(--user-bg) !important;
    border: 1px solid rgba(124,106,247,0.2) !important;
    border-radius: 14px 14px 4px 14px !important;
    color: var(--text-primary) !important;
    font-size: 0.9rem !important;
    padding: 12px 16px !important;
    max-width: 82% !important;
    margin-left: auto !important;
}
#chatbot .message.bot {
    background: var(--bot-bg) !important;
    border: 1px solid var(--border) !important;
    border-radius: 14px 14px 14px 4px !important;
    color: var(--text-primary) !important;
    font-size: 0.9rem !important;
    padding: 12px 16px !important;
    max-width: 88% !important;
    line-height: 1.65 !important;
}
#chatbot code {
    font-family: 'JetBrains Mono', monospace !important;
    background: rgba(124,106,247,0.12) !important;
    color: var(--accent-bright) !important;
    padding: 2px 6px !important;
    border-radius: 5px !important;
    font-size: 0.83em !important;
}
#chatbot pre {
    background: #0d0d14 !important;
    border: 1px solid var(--border-light) !important;
    border-radius: 10px !important;
    padding: 16px !important;
    overflow-x: auto !important;
    margin: 10px 0 !important;
}
#chatbot pre code {
    background: transparent !important;
    color: #c9d1d9 !important;
    padding: 0 !important;
    font-size: 0.85rem !important;
    line-height: 1.6 !important;
}

#msg-input textarea {
    background: var(--bg-tertiary) !important;
    border: 1px solid var(--border) !important;
    border-radius: var(--radius) !important;
    color: var(--text-primary) !important;
    font-family: 'Inter', sans-serif !important;
    font-size: 0.92rem !important;
    padding: 12px 16px !important;
    resize: none !important;
    transition: border-color 0.2s ease !important;
}
#msg-input textarea:focus {
    border-color: var(--accent) !important;
    box-shadow: 0 0 0 3px var(--accent-glow) !important;
}

#send-btn {
    background: linear-gradient(135deg, var(--accent), #5b4fcf) !important;
    color: #fff !important;
    border: none !important;
    border-radius: var(--radius) !important;
    font-family: 'Syne', sans-serif !important;
    font-weight: 600 !important;
    font-size: 0.88rem !important;
    letter-spacing: 0.03em !important;
    padding: 10px 22px !important;
    cursor: pointer !important;
    transition: all 0.2s ease !important;
    box-shadow: 0 4px 15px rgba(124,106,247,0.3) !important;
    height: 100% !important;
}
#send-btn:hover {
    transform: translateY(-1px) !important;
    box-shadow: 0 6px 20px rgba(124,106,247,0.45) !important;
}

label span {
    color: var(--text-secondary) !important;
    font-size: 0.82rem !important;
    font-weight: 500 !important;
    letter-spacing: 0.04em !important;
    text-transform: uppercase !important;
    font-family: 'Syne', sans-serif !important;
}

input[type=range] { accent-color: var(--accent) !important; }

#stats-bar {
    background: var(--bg-tertiary);
    border: 1px solid var(--border);
    border-radius: 8px;
    padding: 8px 14px;
    font-size: 0.78rem;
    color: var(--text-muted);
    font-family: 'JetBrains Mono', monospace;
    letter-spacing: 0.02em;
    margin: 6px 0;
}

#token-info {
    font-family: 'JetBrains Mono', monospace;
    font-size: 0.75rem;
    color: var(--text-muted);
    text-align: right;
    padding: 4px 8px;
}

#system-prompt textarea {
    font-family: 'JetBrains Mono', monospace !important;
    font-size: 0.82rem !important;
    background: var(--bg-tertiary) !important;
    border: 1px solid var(--border) !important;
    color: var(--text-secondary) !important;
    border-radius: var(--radius) !important;
    line-height: 1.6 !important;
}

.status-dot {
    display: inline-block;
    width: 7px; height: 7px;
    border-radius: 50%;
    background: #4ade80;
    box-shadow: 0 0 8px #4ade80;
    animation: pulse 2s ease-in-out infinite;
    margin-right: 6px;
    vertical-align: middle;
}
@keyframes pulse {
    0%, 100% { opacity: 1; }
    50%       { opacity: 0.4; }
}

::-webkit-scrollbar { width: 5px; height: 5px; }
::-webkit-scrollbar-track { background: var(--bg-secondary); }
::-webkit-scrollbar-thumb { background: var(--border-light); border-radius: 10px; }
::-webkit-scrollbar-thumb:hover { background: var(--accent); }
"""

# ── UI ────────────────────────────────────────────────────────────────────────
with gr.Blocks(css=CSS, title="Huihui-Qwen3.5 Chat", theme=gr.themes.Base()) as demo:

    with gr.Group(elem_id="header-block"):
        gr.HTML("""
        <h1>⚡ Huihui-Qwen3.5-2B</h1>
        <p>
            <span class="status-dot"></span>
            abliterated &nbsp;·&nbsp; ZeroGPU &nbsp;·&nbsp; Streaming &nbsp;·&nbsp; Markdown &nbsp;·&nbsp; Export
            &nbsp;|&nbsp;
            <span style="font-family:JetBrains Mono,monospace;font-size:0.8em;color:#7c6af7">
                huihui-ai/Huihui-Qwen3.5-2B-abliterated
            </span>
        </p>
        """)

    with gr.Row(equal_height=False):

        # ── Left: Chat ──────────────────────────────────────────────────
        with gr.Column(scale=7):
            chatbot = gr.Chatbot(
                elem_id="chatbot",
                height=540,
                show_label=False,
            )

            stats_html = gr.HTML('<div id="stats-bar">Henüz mesaj yok.</div>')

            with gr.Row():
                msg = gr.Textbox(
                    elem_id="msg-input",
                    placeholder="Mesajınızı yazın... (Enter = gönder, Shift+Enter = yeni satır)",
                    lines=3,
                    max_lines=8,
                    show_label=False,
                    scale=9,
                )
                send_btn = gr.Button("Gönder ↑", elem_id="send-btn", variant="primary", scale=1)

            with gr.Row():
                clear_btn   = gr.Button("🗑 Sohbeti Temizle", variant="secondary", scale=1)
                token_info  = gr.HTML('<div id="token-info">Tokens: —</div>', scale=1)

        # ── Right: Settings ────────────────────────────────────────────
        with gr.Column(scale=3):
            with gr.Tabs():

                with gr.Tab("🎛 Sistem"):
                    preset_dd = gr.Dropdown(
                        choices=list(SYSTEM_PRESETS.keys()),
                        value="🤖 Varsayılan Asistan",
                        label="Hazır Şablonlar",
                        interactive=True,
                    )
                    system_prompt = gr.Textbox(
                        elem_id="system-prompt",
                        value=SYSTEM_PRESETS["🤖 Varsayılan Asistan"],
                        label="System Prompt",
                        lines=6,
                        placeholder="Modele kimliğini ve davranışını tanımlayın...",
                    )

                with gr.Tab("⚙️ Parametreler"):
                    max_new_tokens = gr.Slider(
                        64, 4096, value=1024, step=64,
                        label="Max Yeni Token",
                        info="Uzun yanıtlar için artırın"
                    )
                    temperature = gr.Slider(
                        0.01, 2.0, value=0.7, step=0.05,
                        label="Temperature",
                        info="Yüksek = yaratıcı, Düşük = tutarlı"
                    )
                    top_p = gr.Slider(
                        0.1, 1.0, value=0.9, step=0.05,
                        label="Top-p (nucleus sampling)",
                    )
                    repetition_penalty = gr.Slider(
                        1.0, 1.5, value=1.1, step=0.02,
                        label="Tekrar Cezası",
                        info="Yüksek = daha az tekrar"
                    )
                    reset_params_btn = gr.Button("↺ Varsayılana Dön", variant="secondary")

                with gr.Tab("📤 Export"):
                    gr.Markdown("Sohbet geçmişini JSON formatında indirin.")
                    export_btn  = gr.Button("💾 JSON İndir", variant="primary")
                    export_file = gr.File(label="İndirme", visible=False)
                    gr.HTML("""
                    <div style="margin-top:16px;font-size:0.78rem;color:#55556a;
                                font-family:JetBrains Mono,monospace;line-height:2;
                                border-top:1px solid #2a2a38;padding-top:12px">
                    <b style="color:#7c6af7">Model</b>&nbsp;&nbsp;&nbsp;Huihui-Qwen3.5-2B-abliterated<br>
                    <b style="color:#7c6af7">Runtime</b>&nbsp;&nbsp;ZeroGPU (A100)<br>
                    <b style="color:#7c6af7">Streaming</b>&nbsp;TextIteratorStreamer<br>
                    <b style="color:#7c6af7">Format</b>&nbsp;&nbsp;&nbsp;ChatML
                    </div>
                    """)

    # ── Event Handlers ─────────────────────────────────────────────────────────

    def user_turn(message, history):
        if not message.strip():
            return "", history
        history = history or []
        history.append({"role": "user", "content": message})
        return "", history

    def bot_turn(history, sys_prompt, max_tok, temp, tp, rep_pen):
        if not history:
            yield history, get_stats([])
            return
        # Son mesaj user mı kontrol et
        last = history[-1]
        last_role = last.get("role") if isinstance(last, dict) else None
        if last_role != "user":
            yield history, get_stats(history)
            return
        user_msg = last.get("content", "") if isinstance(last, dict) else last[0]
        prev_history = history[:-1]
        # Bot placeholder ekle
        history.append({"role": "assistant", "content": ""})
        for partial in generate_stream(user_msg, prev_history, sys_prompt, max_tok, temp, tp, rep_pen):
            history[-1]["content"] = partial
            yield history, get_stats(history)

    # Wiring
    preset_dd.change(update_system_prompt, inputs=preset_dd, outputs=system_prompt)

    msg.submit(user_turn, [msg, chatbot], [msg, chatbot]).then(
        bot_turn,
        [chatbot, system_prompt, max_new_tokens, temperature, top_p, repetition_penalty],
        [chatbot, stats_html]
    )
    send_btn.click(user_turn, [msg, chatbot], [msg, chatbot]).then(
        bot_turn,
        [chatbot, system_prompt, max_new_tokens, temperature, top_p, repetition_penalty],
        [chatbot, stats_html]
    )

    clear_btn.click(
        lambda: ([], '<div id="stats-bar">Henüz mesaj yok.</div>'),
        outputs=[chatbot, stats_html]
    )

    msg.change(update_token_count, inputs=msg, outputs=token_info)

    reset_params_btn.click(
        reset_parameters,
        outputs=[max_new_tokens, temperature, top_p, repetition_penalty]
    )

    export_btn.click(
        export_chat, inputs=[chatbot, system_prompt], outputs=export_file
    ).then(lambda: gr.update(visible=True), outputs=export_file)


demo.queue(max_size=10).launch()