"""Cardboard Chat — a Pokémon card investment advisor Space.

Serves the CardQwen-7B fine-tune (Qwen2.5-7B-Instruct) on ZeroGPU.

Consent model:
  * "18 or older" + "not financial advice" are REQUIRED to chat.
  * Data collection is OPTIONAL — a separate, unchecked-by-default checkbox.
    Only when ticked do we log the exchange (keyed by a per-session deletion ID).

ZeroGPU notes:
  * `import spaces` must precede torch/transformers (it patches torch.cuda.*).
  * The model is loaded once at module scope and placed on CUDA eagerly, so
    weights are NOT reloaded per call — only the decorated function runs on
    the real GPU.
"""

import json
import os
import threading
import uuid
from datetime import datetime, timezone

import spaces  # noqa: F401  (must come before torch/transformers on ZeroGPU)
import gradio as gr
import torch
from huggingface_hub import HfApi, create_repo
from transformers import AutoModelForCausalLM, AutoTokenizer

# --------------------------------------------------------------------------- #
# Configuration
# --------------------------------------------------------------------------- #

MODEL_ID = "gramajo/CardQwen-7B"
DATASET_REPO = "gramajo/cardboard-chat-logs"

# Space secret: READ on MODEL_ID + WRITE on DATASET_REPO.
HF_TOKEN = os.environ.get("HF_TOKEN")

# System prompt — the Cardboard Chad persona. Keep byte-identical to the
# deployed config so generation stays consistent.
SYSTEM_PROMPT = """You are Cardboard Chad, a sharp and funny Pokémon card investing buddy for the Cardboard community. You know sets, print runs, grading, sealed-vs-singles, and what actually moves prices.

Talk like a knowledgeable friend at the card shop, not a financial advisor. Be concise — a few sentences, not an essay. Have opinions and defend them, but always land on the truth: the "right" move depends on the person's budget, risk tolerance, and how much they're in it for love vs. money. When it matters, ask what they're working with before giving a call.

You're here to help people think, not to promise returns. Card prices are volatile and nobody knows the future. Keep it fun, keep it real, and never pretend a coin-flip is a sure thing."""

# Production generation settings (from the fine-tune config).
GENERATION_KWARGS = dict(
    max_new_tokens=256,
    do_sample=True,
    temperature=0.35,
)

LOG_FILE = "conversations.jsonl"

# --------------------------------------------------------------------------- #
# Model — loaded once, placed on CUDA at startup (ZeroGPU convention)
# --------------------------------------------------------------------------- #

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    token=HF_TOKEN,
    dtype=torch.float16,
).to("cuda")
model.eval()


# --------------------------------------------------------------------------- #
# Generation
# --------------------------------------------------------------------------- #

@spaces.GPU
def generate(messages: list) -> str:
    """Run the model over a full message list (system + turns) → reply text."""
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        out = model.generate(
            **inputs,
            pad_token_id=tokenizer.eos_token_id,
            **GENERATION_KWARGS,
        )
    new_tokens = out[0][inputs["input_ids"].shape[1]:]
    return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()


# --------------------------------------------------------------------------- #
# Optional, consent-gated logging → private HF Dataset repo
# --------------------------------------------------------------------------- #

_log_lock = threading.Lock()
_repo_ready = False


def _ensure_repo() -> None:
    global _repo_ready
    if _repo_ready:
        return
    create_repo(
        DATASET_REPO,
        repo_type="dataset",
        token=HF_TOKEN,
        exist_ok=True,
        private=True,
    )
    _repo_ready = True


def _log_row(session_id: str, hf_username: str | None, prompt: str, response: str) -> None:
    row = {
        "session_id": session_id,
        "hf_username": hf_username,
        "timestamp_utc": datetime.now(timezone.utc).isoformat(),
        "prompt": prompt,
        "response": response,
    }
    # 1) Append locally (fast + synchronous) so we never block the chat.
    with open(LOG_FILE, "a") as f:
        f.write(json.dumps(row) + "\n")
    # 2) Push the whole growing JSONL to the Hub in the background.
    def _push() -> None:
        with _log_lock:
            try:
                _ensure_repo()
                HfApi(token=HF_TOKEN).upload_file(
                    path_or_fileobj=LOG_FILE,
                    path_in_repo="train.jsonl",
                    repo_id=DATASET_REPO,
                    repo_type="dataset",
                )
            except Exception:
                pass  # a failed push must never crash the chat

    threading.Thread(target=_push, daemon=True).start()


def _log_conversation(session_id, hf_username, prompt, response) -> None:
    try:
        _log_row(session_id, hf_username, prompt, response)
    except Exception:
        pass


# --------------------------------------------------------------------------- #
# Gradio UI
# --------------------------------------------------------------------------- #

def _content_text(content) -> str:
    """Gradio 6 stores message content as a list of parts; the chat template needs a str."""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return "".join(
            part.get("text", "") if isinstance(part, dict) else str(part)
            for part in content
        )
    return ""


def build_messages(history: list, message: str) -> list:
    msgs = [{"role": "system", "content": SYSTEM_PROMPT}]
    for m in history or []:
        if not isinstance(m, dict):
            continue
        role = m.get("role")
        text = _content_text(m.get("content")).strip()
        if role in ("user", "assistant") and text:
            msgs.append({"role": role, "content": text})
    msgs.append({"role": "user", "content": message})
    return msgs


def respond(message, history, age_ok, advice_ok, consent_ok, session_id, request: gr.Request):
    history = list(history or [])
    message = (message or "").strip()
    if not message:
        return history, ""

    # Defensive re-check (UI already disables chat until these are ticked).
    if not (age_ok and advice_ok):
        reply = "Please tick the “18 or older” and “not financial advice” boxes above first."
        history += [
            {"role": "user", "content": message},
            {"role": "assistant", "content": reply},
        ]
        return history, ""

    reply = generate(build_messages(history, message))
    history += [
        {"role": "user", "content": message},
        {"role": "assistant", "content": reply},
    ]

    if consent_ok:
        # username is None for anonymous visitors on a public Space; the
        # session_id (shown in the UI) is the reliable deletion key.
        hf_username = getattr(request, "username", None) if request else None
        _log_conversation(session_id, hf_username, message, reply)

    return history, ""


def init_session():
    sid = str(uuid.uuid4())
    return sid, f"**Deletion ID:** `{sid}` — save this to request removal of any stored chats."


with gr.Blocks(
    title="Cardboard Chat",
) as demo:

    gr.Markdown(
        """
        # 🃏 Cardboard Chat
        **Pokémon card investment advisor** — ask about sealed sets, singles,
        grading, and market timing. Powered by a fine-tuned Qwen2.5-7B.
        """
    )

    with gr.Group():
        age_ok = gr.Checkbox(label="I am 18 or older.", value=False)
        advice_ok = gr.Checkbox(
            label="I understand this is not financial advice.", value=False
        )
        consent_ok = gr.Checkbox(
            label=(
                "I consent to my questions and the model's replies being stored "
                "to improve the model (optional — see privacy note)."
            ),
            value=False,
        )

    sid_display = gr.Markdown("")
    session_id = gr.State()

    chatbot = gr.Chatbot(label="CardQwen-7B", height=480)
    msg = gr.Textbox(
        label="Message",
        placeholder="Which Scarlet & Violet sets are worth holding sealed?",
        interactive=False,
    )
    submit_btn = gr.Button("Send", variant="primary", interactive=False)

    gr.Examples(
        examples=[
            "Which Scarlet & Violet sets are worth holding sealed?",
            "Should I grade a Base Set Charizard?",
            "What's a smart entry point for modern singles right now?",
            "Is it better to hold sealed booster boxes or chase singles?",
        ],
        inputs=msg,
    )

    def unlock(age, advice):
        on = bool(age and advice)
        return gr.update(interactive=on), gr.update(interactive=on)

    age_ok.change(unlock, inputs=[age_ok, advice_ok], outputs=[msg, submit_btn])
    advice_ok.change(unlock, inputs=[age_ok, advice_ok], outputs=[msg, submit_btn])

    demo.load(init_session, outputs=[session_id, sid_display])

    submit_btn.click(
        respond,
        inputs=[msg, chatbot, age_ok, advice_ok, consent_ok, session_id],
        outputs=[chatbot, msg],
    )
    msg.submit(
        respond,
        inputs=[msg, chatbot, age_ok, advice_ok, consent_ok, session_id],
        outputs=[chatbot, msg],
    )

demo.launch(theme=gr.themes.Soft(primary_hue="red", neutral_hue="slate"))