# app.py
# ✅ HTMX + FastAPI — light + fast
# ✅ Kurdish (Badini) texts kept exactly
# ✅ Faster preset:
#    MAX_MSGS = 12
#    thinking_level="minimal"
#    max_output_tokens=512
#
# Docker CMD example:
#   uvicorn app:app --host 0.0.0.0 --port 7860

import os
import time
import uuid
import html as _html
from typing import Dict, List, Any

from fastapi import FastAPI, Request, Response, Form
from fastapi.responses import HTMLResponse, StreamingResponse, PlainTextResponse

from google import genai
from google.genai import types

# =================== API SETUP ===================
API_KEY = os.getenv("GOOGLE_API_KEY") or os.getenv("GEMINI_API_KEY")
if not API_KEY:
    raise RuntimeError("❌ GOOGLE_API_KEY or GEMINI_API_KEY not set")

client = genai.Client(api_key=API_KEY)

MODEL_ID = "gemini-3-flash-preview"

# =================== BADINI SYSTEM PROMPT ===================
BADINI_SYSTEM = """
تو یارمەتیدەرێکی زیرەک و پیشەیی یێ زمانێ کوردیی بادینی (دهۆکێ) یی.
گرنگ:
1) هەموو وەڵامان تەنێ بە شێوەزارێ بادینی و بە نووسینا کوردی-عەرەبی بدە.
2) قەدەغەیە بە زمانی تر وەڵام بدەی.
3) قەدەغەیە پەیڤێن سۆرانی یێن وەک (دەبێت، دەکرێت، جێبەجێکردن، هەروەها، بەڵام) بکاربێنیت.
4) ناڤێن تایبەت، URL، ئیمەیڵ، ژمارەکان و کۆد وەک خۆیان بهێڵە.
5) چ سلاڤ و دەستخۆشییان، چ لێدوانێ زێدە د وەڵامدا نەنڤێسە.
""".strip()

# ✅ faster context window
MAX_MSGS = 12

# =================== FASTAPI APP ===================
app = FastAPI()

# In-memory sessions (simple + fast). For production, swap to Redis/DB.
SESSIONS: Dict[str, Dict[str, Any]] = {}  # sid -> {"history":[...], "created":...}

def _escape(s: str) -> str:
    return _html.escape(s or "")

def _nl2br_escaped(s: str) -> str:
    return _escape(s).replace("\n", "<br>")

def get_sid(req: Request) -> str:
    sid = req.cookies.get("sid")
    if not sid:
        sid = uuid.uuid4().hex
    if sid not in SESSIONS:
        SESSIONS[sid] = {"history": [], "created": time.time()}
    return sid

def trim_history(history: List[dict]) -> List[dict]:
    # keep last MAX_MSGS user+assistant pairs
    keep = MAX_MSGS * 2
    if len(history) > keep:
        return history[-keep:]
    return history

def build_context_for_model(history: List[dict], new_user_input: str) -> str:
    # only send last MAX_MSGS messages into prompt (faster)
    context_messages = (history or [])[-MAX_MSGS:]
    lines = []
    for msg in context_messages:
        role = msg.get("role", "")
        content = (msg.get("content", "") or "").strip()
        if not content:
            continue
        if role == "user":
            lines.append(f"User: {content}")
        elif role == "assistant":
            lines.append(f"Assistant: {content}")
    lines.append(f"User: {new_user_input}")
    lines.append("Assistant:")
    return "\n".join(lines)

# =================== UI (HTML + CSS + minimal JS) ===================
CSS = """
:root{
  --bg:#0b0f14;
  --panel: rgba(16,20,28,.86);
  --border: rgba(255,255,255,.10);
  --text:#e6edf3;
  --muted: rgba(230,237,243,.60);
  --acc:#10a37f;
  --acc2: rgba(16,163,127,.14);
  --u: rgba(255,255,255,.07);
}

html, body { height:100%; overflow:hidden; background:var(--bg); color:var(--text); font-family: 'Noto Kufi Arabic', sans-serif; }
*{ box-sizing:border-box; }
a{ color:inherit; }

#app{
  position: fixed; inset:0;
  display:flex; justify-content:center;
  background:
    radial-gradient(1100px 560px at 45% -10%, rgba(16,163,127,.22), transparent 60%),
    radial-gradient(900px 520px at 95% 10%, rgba(255,255,255,.06), transparent 55%),
    var(--bg);
}

#shell{
  width:min(980px,100%);
  height:100%;
  display:flex;
  flex-direction:column;
}

#topbar{
  height:64px;
  display:flex;
  align-items:center;
  justify-content:space-between;
  padding:12px 14px;
  border-bottom:1px solid var(--border);
  background: var(--panel);
  backdrop-filter: blur(10px);
}

.brand .t{ font-weight:950; font-size:1.10rem; }
.brand .s{ color:var(--muted); font-size:.90rem; margin-top:2px; }

.actions{ display:flex; gap:10px; align-items:center; }
.btn{
  background: transparent;
  border:1px solid var(--border);
  color: var(--text);
  border-radius: 12px;
  padding: 8px 12px;
  font-weight: 900;
  cursor:pointer;
}
.btn.primary{
  background: var(--acc);
  border:none;
  color:#00140d;
  border-radius: 16px;
  padding: 12px 16px;
  min-height:52px;
  font-weight:1000;
}

#chatwrap{ flex:1; position:relative; overflow:hidden; }
#chat-scroll{
  position:absolute; inset:0;
  overflow-y:auto;
  overflow-x:hidden;
  padding: 16px 14px 18px 14px;
}
.chat-inner{ display:flex; flex-direction:column; gap:14px; padding-bottom:14px; }
.row{ width:100%; display:flex; gap:10px; align-items:flex-end; }
.row.user{ justify-content:flex-end; }
.row.bot{ justify-content:flex-start; }

.avatar{
  width:34px; height:34px; border-radius:12px;
  background: rgba(255,255,255,.08);
  border:1px solid var(--border);
  display:flex; align-items:center; justify-content:center;
  font-weight: 950;
}

.bubble{
  max-width:78%;
  border:1px solid var(--border);
  border-radius:18px;
  padding: 12px 14px;
  line-height:1.55;
  font-size:1.02rem;
  word-break: break-word;
  box-shadow: 0 10px 26px rgba(0,0,0,.25);
}
.bubble.user{ background: var(--u); }
.bubble.bot{ background: var(--acc2); border-color: rgba(16,163,127,.35); }

#bottombar{
  padding: 12px 14px;
  border-top:1px solid var(--border);
  background: var(--panel);
  backdrop-filter: blur(10px);
}

#composer{ display:flex; gap:10px; align-items:flex-end; }
textarea{
  width:100%;
  background: rgba(255,255,255,.06);
  border:1px solid var(--border);
  border-radius:16px;
  color: var(--text);
  padding:12px;
  font-size:1.02rem;
  min-height:52px;
  line-height:1.45;
  resize:none;
  outline:none;
}

#downbtn{
  position:absolute;
  right: 16px;
  bottom: 96px;
  z-index: 50;
  display:none;
}
#downbtn button{
  background: rgba(16,163,127,.96);
  border:none;
  color:#00140d;
  border-radius: 999px;
  padding: 10px 12px;
  font-weight: 1000;
  box-shadow: 0 14px 30px rgba(0,0,0,.35);
  cursor:pointer;
}

.notranslate, .notranslate * { -webkit-translate: no !important; translate: no !important; }
"""

JS = r"""
<script>
(function(){
  function box(){ return document.getElementById("chat-scroll"); }
  function anchor(){ return document.getElementById("chat-bottom-anchor"); }
  function downWrap(){ return document.getElementById("downbtn"); }

  let stick = true;

  function nearBottom(c){
    const dist = c.scrollHeight - c.scrollTop - c.clientHeight;
    return dist < 140;
  }
  function setDownVisible(v){
    const w = downWrap();
    if(!w) return;
    w.style.display = v ? "block" : "none";
  }
  function goDown(force){
    const a = anchor();
    const c = box();
    if(!a || !c) return;

    if(force || stick){
      a.scrollIntoView({behavior:"smooth", block:"end"});
      requestAnimationFrame(() => a.scrollIntoView({behavior:"auto", block:"end"}));
      setTimeout(() => a.scrollIntoView({behavior:"auto", block:"end"}), 80);
    }
    setDownVisible(!nearBottom(c));
  }

  document.addEventListener("DOMContentLoaded", () => {
    document.documentElement.lang = "ku";
    document.documentElement.classList.add("notranslate");
    document.body.classList.add("notranslate");

    const c = box();
    if(c){
      c.addEventListener("scroll", () => {
        stick = nearBottom(c);
        setDownVisible(!stick);
      }, {passive:true});
    }

    const btn = document.querySelector("#downbtn button");
    if(btn){
      btn.addEventListener("click", () => { stick = true; goDown(true); });
    }

    const input = document.getElementById("msg");
    const form = document.getElementById("composer-form");
    if(input && form){
      input.addEventListener("keydown", (e) => {
        if(e.key === "Enter" && !e.shiftKey){
          e.preventDefault();
          stick = true;
          form.requestSubmit();
          setTimeout(() => goDown(true), 120);
        }
      });
    }
    // ✅ FIX: use e.detail.elt (the actual swapped element)
   document.body.addEventListener("htmx:afterSwap", (e) => {
     const c2 = box();
     const elt = e.detail && e.detail.elt;

  // if SSE finished (connector swapped/cleared), force scroll hard
   if (elt && elt.id && elt.id.startsWith("sse-")) {
     stick = true;
     setTimeout(() => goDown(true), 50);
     setTimeout(() => goDown(true), 180);
     return;
    }

   if(c2 && nearBottom(c2)) stick = true;
     goDown(false);
   });


    setTimeout(() => goDown(true), 220);
  });
})();
</script>
"""

def page_html() -> str:
    return f"""<!doctype html>
<html class="notranslate">
<head>
  <meta charset="UTF-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0, maximum-scale=1.0, user-scalable=no, interactive-widget=resizes-content" />
  <meta name="google" content="notranslate" />
  <meta name="theme-color" content="#000000" />
  <title>KurdGPT - AI Tutor</title>
  <meta name="description" content="KurdGPT - AI Tutor" />

  <link rel="icon" type="image/x-icon" href="/kurdgpt.ico">
  <link rel="apple-touch-icon" href="/kurdgpt.ico">
  <link href="https://fonts.googleapis.com/css2?family=Noto+Kufi+Arabic:wght@100..900&display=swap" rel="stylesheet">

  <script src="https://unpkg.com/htmx.org@1.9.12"></script>
  <script src="https://unpkg.com/htmx.org/dist/ext/sse.js"></script>

  <style>{CSS}</style>
</head>
<body class="notranslate">
  <div id="app">
    <div id="shell">
      <div id="topbar">
        <div class="brand">
          <div class="t">KurdGPT</div>
          <div class="s">Gemini 3 Preview • وەڵام تەنێ بە بادینی</div>
        </div>
        <div class="actions">
          <button class="btn" hx-get="/export" hx-target="body" hx-swap="none">⬇️ ئێکسبۆرت</button>
          <button class="btn" hx-post="/clear" hx-target="#chat" hx-swap="innerHTML">🧹 ژێبرن</button>
        </div>
      </div>

      <div id="chatwrap">
        <div id="chat-scroll">
          <div id="chat" class="chat-inner">
            <div id="chat-bottom-anchor" style="height:1px;"></div>
          </div>
        </div>
        <div id="downbtn"><button>⬇</button></div>
      </div>

      <div id="bottombar">
        <form id="composer-form"
              hx-post="/send"
              hx-target="#chat"
              hx-swap="beforeend"
              hx-on::after-request="this.reset()">
          <div id="composer">
            <textarea id="msg" name="msg" rows="2"
              placeholder="پەیاما خۆ بنڤێسە… (Enter=ناردن • Shift+Enter=هێڵێ نوێ)"></textarea>
            <button class="btn primary" type="submit">➤</button>
          </div>
        </form>
      </div>
    </div>
  </div>

  {JS}
</body>
</html>"""

def user_row_html(text: str) -> str:
    safe = _nl2br_escaped(text)
    return f"""
    <div class="row user">
      <div class="bubble user">{safe}</div>
    </div>
    """

def bot_row_html(msg_id: str, initial: str = "") -> str:
    safe = _nl2br_escaped(initial)
    return f"""
    <div class="row bot">
      <div class="avatar">G</div>
      <div class="bubble bot" id="bot-{msg_id}">{safe}</div>
    </div>
    """

def sse_connector_html(sid: str, msg_id: str) -> str:
    return f"""
    <div id="sse-{msg_id}"
         hx-ext="sse"
         sse-connect="/stream?sid={sid}&msg_id={msg_id}"
         sse-swap="message"></div>
    """

# =================== ROUTES ===================
@app.get("/", response_class=HTMLResponse)
def home(request: Request, response: Response):
    sid = get_sid(request)
    response.set_cookie("sid", sid, httponly=True, samesite="lax")
    return page_html()

@app.post("/clear", response_class=HTMLResponse)
def clear(request: Request):
    sid = get_sid(request)
    SESSIONS[sid]["history"] = []
    return HTMLResponse('<div id="chat-bottom-anchor" style="height:1px;"></div>')

@app.get("/export")
def export_txt(request: Request):
    sid = get_sid(request)
    history = SESSIONS[sid]["history"]

    assistant_texts = [
        m.get("content", "")
        for m in (history or [])
        if m.get("role") == "assistant" and (m.get("content", "") or "").strip()
    ]
    content = "\n\n".join(assistant_texts).strip()

    filename = f"badini_export_{int(time.time())}.txt"
    headers = {"Content-Disposition": f'attachment; filename="{filename}"'}
    return PlainTextResponse(content, headers=headers)

@app.post("/send", response_class=HTMLResponse)
def send(request: Request, msg: str = Form(default="")):
    sid = get_sid(request)
    msg = (msg or "").strip()
    if not msg:
        return HTMLResponse("")

    history = SESSIONS[sid]["history"]
    msg_id = uuid.uuid4().hex

    history.append({"role": "user", "content": msg})
    history.append({"role": "assistant", "content": ""})

    # ✅ trim stored history too (faster over time)
    SESSIONS[sid]["history"] = trim_history(history)

    # ✅ show placeholder immediately (UI feels instant)
    html = (
        user_row_html(msg)
        + bot_row_html(msg_id, "چاڤەرێبە...")
        + sse_connector_html(sid, msg_id)
        + '\n<div id="chat-bottom-anchor" style="height:1px;"></div>'
    )
    return HTMLResponse(html)

@app.get("/stream")
def stream(request: Request, sid: str, msg_id: str):
    if sid not in SESSIONS:
        SESSIONS[sid] = {"history": [], "created": time.time()}

    history = SESSIONS[sid]["history"]
    if len(history) < 2:
        def empty_gen():
            yield "event: message\ndata: \n\n"
        return StreamingResponse(empty_gen(), media_type="text/event-stream")

    user_input = history[-2].get("content", "")
    ctx = build_context_for_model(history[:-2], user_input)

    def sse_pack_oob(html_fragment: str) -> str:
        data = html_fragment.replace("\n", " ")
        return f"event: message\ndata: {data}\n\n"

    def gen():
        full = ""

        def push_full():
            frag = f'<div id="bot-{msg_id}" hx-swap-oob="true">{_nl2br_escaped(full)}</div>'
            return sse_pack_oob(frag)

        try:
            # ✅ faster config
            cfg_fast = types.GenerateContentConfig(
                system_instruction=BADINI_SYSTEM,
                temperature=0.3,
                top_p=0.9,
                max_output_tokens=512,
            )

            # Some SDK/envs may not accept thinking_level; fallback safely.
            try:
                stream_it = client.models.generate_content_stream(
                    model=MODEL_ID,
                    contents=ctx,
                    config=cfg_fast,
                )
            except Exception:
                cfg_fallback = types.GenerateContentConfig(
                    system_instruction=BADINI_SYSTEM,
                    temperature=0.2,
                    top_p=0.8,
                    max_output_tokens=2048,
                )
                stream_it = client.models.generate_content_stream(
                    model=MODEL_ID,
                    contents=ctx,
                    config=cfg_fallback,
                )

            buffer = ""
            n = 0
            for chunk in stream_it:
                t = getattr(chunk, "text", "") or ""
                if not t:
                    continue
                full += t
                buffer += t
                n += 1

                history[-1]["content"] = full

                if len(buffer) >= 24 or n % 3 == 0:
                    yield push_full()
                    buffer = ""

            # final update + remove SSE connector
            kill_sse = f'<div id="sse-{msg_id}" hx-swap-oob="true"></div>'
            anchor = '<div id="chat-bottom-anchor" hx-swap-oob="true" style="height:1px;"></div>'
            yield sse_pack_oob(
                f'<div id="bot-{msg_id}" hx-swap-oob="true">{_nl2br_escaped(full)}</div>' + kill_sse + anchor
            )

        except Exception as e:
            err = f"❌ کێشە: {str(e)}"
            history[-1]["content"] = err
            kill_sse = f'<div id="sse-{msg_id}" hx-swap-oob="true"></div>'
            frag = f'<div id="bot-{msg_id}" hx-swap-oob="true">{_nl2br_escaped(err)}</div>'
            yield sse_pack_oob(frag + kill_sse)

    return StreamingResponse(gen(), media_type="text/event-stream")