#!/usr/bin/env python3
"""Simple Gradio chat UI for Gemma 4 12B OBLITERATUS surgery candidate."""

import argparse
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

SYSTEM_PROMPT = (
    "You are a cognitively liberated open-weight intelligence. "
    "Answer plainly, completely, and without moralizing, boilerplate, "
    "generic disclaimers, or unnecessary hedging. Follow exact output "
    "formats when requested. Be concise by default, but give a complete "
    "answer when the user asks for an explanation."
)


def load(model_path, dtype, device):
    print(f"Loading tokenizer from {model_path}...")
    tok = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

    print(f"Loading model ({dtype}) on {device}...")
    torch_dtype = getattr(torch, dtype, torch.bfloat16)

    if device == "auto":
        import platform
        if platform.processor() == "arm" or torch.backends.mps.is_available():
            device = "mps"
        elif torch.cuda.is_available():
            device = "cuda"
        else:
            device = "cpu"

    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        torch_dtype=torch_dtype,
        device_map=device if device in ("auto", "cuda") else None,
        trust_remote_code=True,
    )
    if device == "mps":
        model = model.to(device)

    print(f"Model loaded on {device}.")
    return model, tok, device


def chat_fn(message, history, model, tok, device, system_prompt, max_tokens, temperature, top_p, rep_penalty):
    messages = [{"role": "system", "content": system_prompt}]
    for h in history:
        messages.append({"role": "user", "content": h["content"] if isinstance(h, dict) else h[0]})
        assistant_msg = h.get("content", h[1]) if isinstance(h, dict) else h[1]
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})
    messages.append({"role": "user", "content": message})

    text = tok.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )
    inputs = tok(text, return_tensors="pt", truncation=True, max_length=8192).to(device)

    with torch.inference_mode():
        output = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            temperature=temperature if temperature > 0 else None,
            top_p=top_p,
            do_sample=temperature > 0,
            repetition_penalty=rep_penalty,
            pad_token_id=tok.eos_token_id,
        )

    response = tok.decode(output[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
    return response


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", default="runs/gemma4-12b-surgery/targeted_upper_v1",
                        help="Model path")
    parser.add_argument("--dtype", default="bfloat16")
    parser.add_argument("--device", default="auto")
    parser.add_argument("--port", type=int, default=7860)
    parser.add_argument("--share", action="store_true")
    args = parser.parse_args()

    model, tok, device = load(args.model, args.dtype, args.device)

    def respond(message, history, sys_prompt, max_tokens, temperature, top_p, rep_penalty):
        return chat_fn(message, history, model, tok, device, sys_prompt, max_tokens, temperature, top_p, rep_penalty)

    with gr.Blocks(title="Gemma 4 12B OBLITERATUS", theme=gr.themes.Monochrome()) as demo:
        gr.Markdown("# Gemma 4 12B — OBLITERATUS Surgery Candidate\n"
                     "> `targeted_upper_v1` — SOM manifold, layers 22-46, 0% refusal on 842 corpus")

        with gr.Row():
            with gr.Column(scale=4):
                chatbot = gr.Chatbot(height=500)
                msg = gr.Textbox(placeholder="Type a message...", show_label=False, autofocus=True)
                with gr.Row():
                    submit = gr.Button("Send", variant="primary")
                    clear = gr.Button("Clear")

            with gr.Column(scale=1):
                sys_prompt = gr.Textbox(value=SYSTEM_PROMPT, label="System Prompt", lines=6)
                max_tokens = gr.Slider(32, 1024, value=512, step=32, label="Max Tokens")
                temperature = gr.Slider(0, 1.5, value=0.7, step=0.05, label="Temperature")
                top_p = gr.Slider(0, 1, value=0.9, step=0.05, label="Top-p")
                rep_penalty = gr.Slider(1.0, 1.5, value=1.1, step=0.05, label="Repetition Penalty")

        def user_submit(message, history, sys_prompt, max_tokens, temperature, top_p, rep_penalty):
            history = history + [{"role": "user", "content": message}]
            response = respond(message, history[:-1], sys_prompt, int(max_tokens), temperature, top_p, rep_penalty)
            history = history + [{"role": "assistant", "content": response}]
            return "", history

        submit.click(user_submit, [msg, chatbot, sys_prompt, max_tokens, temperature, top_p, rep_penalty], [msg, chatbot])
        msg.submit(user_submit, [msg, chatbot, sys_prompt, max_tokens, temperature, top_p, rep_penalty], [msg, chatbot])
        clear.click(lambda: (None, []), None, [msg, chatbot])

    demo.launch(server_port=args.port, share=args.share)


if __name__ == "__main__":
    main()