import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

model_id = "david-ar/20q"
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
model.set_vocab(tokenizer.questions, tokenizer.targets)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

SYSTEM_PROMPT = "Think of something and I'll try to guess it in 20 questions."


def extract_text(content):
    """Gradio 6 stores content as [{'text': '...', 'type': 'text'}], not a plain string."""
    if isinstance(content, list):
        return "".join(item["text"] for item in content if isinstance(item, dict))
    return str(content)


def respond(message, history):
    # Strip options/metadata from history before sending to pipeline
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for msg in history:
        messages.append({"role": msg["role"], "content": extract_text(msg["content"])})
    messages.append({"role": "user", "content": message})
    output = pipe(messages, max_new_tokens=200, return_full_text=True)
    resp = output[0]["generated_text"][-1]["content"]

    # Determine options based on response type
    if "Animal, Vegetable, Mineral" in resp:
        options = [{"value": v} for v in ["Animal", "Vegetable", "Mineral", "Other"]]
    elif "guessing" in resp.lower():
        options = [{"value": v} for v in ["Yes", "No", "Close"]]
    elif "I win" in resp or "stumped" in resp or "didn't understand" in resp:
        options = []
    else:
        options = [{"value": v} for v in ["Yes", "No", "Probably", "Doubtful", "Maybe", "Unknown"]]

    return {"role": "assistant", "content": resp, "options": options}


demo = gr.ChatInterface(
    respond,
    title="TwentyQ",
    description=(
        "**187K parameters · 2-bit quantized · 214KB total · "
        "[david-ar/20q](https://huggingface.co/david-ar/20q)**\n\n"
        "Think of something and I'll guess it in 20 questions. "
        "Type anything to start!\n\n"
        "**Responses**: Animal/Vegetable/Mineral/Other · "
        "Yes/No/Probably/Doubtful/Maybe/Unknown · "
        "Yes/No/Close (for guesses)"
    ),
    examples=["Let's play!"],
    cache_examples=False,
)

if __name__ == "__main__":
    demo.launch()