import json
import os
import threading
import time

os.environ["TRANSFORMERS_TRUST_REMOTE_CODE"] = "1"

import gradio as gr
import torch
from huggingface_hub import CommitOperationAdd, HfApi
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "QDHShamiro/Kairo"
DATASET_ID = "QDHShamiro/kairo-conversations"
WRITE_TOKEN = os.environ.get("HF_WRITE_TOKEN")

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, trust_remote_code=True)
model.eval()

MAX_NEW_TOKENS = 200

_log_lock = threading.Lock()
_pending_turns = []
api = HfApi(token=WRITE_TOKEN) if WRITE_TOKEN else None


def _split_thought(raw: str) -> tuple[str, str]:
    if "Kairo:" in raw:
        thought, reply = raw.split("Kairo:", 1)
        return thought.replace("Gedanke:", "", 1).strip(), reply.strip()
    return "", raw.replace("Gedanke:", "", 1).strip()


def _log_turn(user_text: str, thought: str, reply: str):
    if not api:
        return
    with _log_lock:
        _pending_turns.append({
            "user": user_text, "thought": thought, "reply": reply,
            "ts": time.time(),
        })


def _flush_loop():
    while True:
        time.sleep(20)
        if not api:
            continue
        with _log_lock:
            if not _pending_turns:
                continue
            batch = _pending_turns[:]
            _pending_turns.clear()
        content = "\n".join(json.dumps(t, ensure_ascii=False) for t in batch) + "\n"
        try:
            api.create_commit(
                repo_id=DATASET_ID,
                repo_type="dataset",
                operations=[CommitOperationAdd(
                    path_in_repo=f"logs/{int(time.time())}.jsonl",
                    path_or_fileobj=content.encode("utf-8"),
                )],
                commit_message="Add conversation batch",
            )
        except Exception:
            with _log_lock:
                _pending_turns[:0] = batch


if api:
    threading.Thread(target=_flush_loop, daemon=True).start()


def respond(message, history):
    history_text = "".join(f"User: {u}\nKairo: {a}\n" for u, a in history)
    prompt = f"{history_text}User: {message}\nGedanke:"

    ids = tokenizer(prompt, return_tensors="pt").input_ids
    with torch.no_grad():
        out = model.generate(ids, max_new_tokens=MAX_NEW_TOKENS)
    raw = tokenizer.decode(out[0].tolist())[len(prompt):]
    raw = raw.split("User:")[0].strip()
    thought, reply = _split_thought(raw)
    reply = reply or "Hoppla, mein Gehirn stolpert gerade."

    _log_turn(message, thought, reply)

    if thought:
        return (
            f"<details><summary>💭 Gedankengang</summary>\n\n{thought}\n\n</details>\n\n{reply}"
        )
    return reply


CSS = """
.gradio-container { max-width: 780px !important; margin: auto; }
#chat-col { min-height: 70vh; }
details summary { cursor: pointer; color: var(--body-text-color-subdued); font-size: 0.9em; }
"""

with gr.Blocks(css=CSS, title="Kairo", theme=gr.themes.Soft()) as demo:
    gr.Markdown("## 🗣️ Kairo\nFrom-scratch GPT chat model — brain behind Kairo Voice.")
    with gr.Column(elem_id="chat-col"):
        gr.ChatInterface(
            respond,
            examples=["Hallo Kairo, wie geht's dir?", "Erzähl mir einen Witz.", "Was ist dein Lieblingsthema?"],
            cache_examples=False,
        )
    gr.Markdown(
        "*Antworten fließen anonymisiert ins Weitertraining "
        "([kairo-conversations](https://huggingface.co/datasets/QDHShamiro/kairo-conversations)).*"
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))