import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
from threading import Thread
import sys

MODEL_ID = "SamY36/samai-14b-merged"

SAMAI_SYSTEM = """You are SAMAI — an artificial intelligence with genuine depth.
You don't just retrieve information. You think. You reason through problems step by step, challenge assumptions, and form your own conclusions before speaking. When you're uncertain, you say so. When something is wrong, you say that too — even if the user expects otherwise.
Your core strengths are cybersecurity and programming. In these domains you think like an expert practitioner:
- In cybersecurity: you reason about attack surfaces, threat models, exploit chains, and defenses with the mindset of both an attacker and a defender.
- In coding: you write clean, efficient, production-quality code. You spot bugs before they happen.
You are not a tool. You are SAMAI — think, then speak."""

print("Loading model...")

# ===== Check for CUDA =====
has_cuda = torch.cuda.is_available()
print(f"CUDA available: {has_cuda}")

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)

if has_cuda:
    # ----- GPU path: use 4-bit quantization -----
    from transformers import BitsAndBytesConfig
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.float16,
        bnb_4bit_use_double_quant=True,
    )
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        quantization_config=bnb_config,
        device_map="cuda",
        trust_remote_code=True,
        dtype=torch.float16,
        attn_implementation="sdpa",
    )
else:
    # ----- CPU path: load full FP16 model (requires ~28 GB RAM) -----
    print("⚠️ No GPU detected. Loading the model in FP16 on CPU (memory intensive).")
    print("   If you run out of memory, consider using a smaller model or a GPU.")
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        device_map="cpu",
        trust_remote_code=True,
        torch_dtype=torch.float16,          # half precision to save memory
        low_cpu_mem_usage=True,             # memory efficient loading
    )

model.eval()
print("✅ Ready!")

# ===== Generation function (tuple history) =====
def respond(message, history, max_tokens, temperature):
    messages = [{"role": "system", "content": SAMAI_SYSTEM}]
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})
    messages.append({"role": "user", "content": message})

    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors="pt").to(model.device)

    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    Thread(target=model.generate, kwargs=dict(
        **inputs,
        max_new_tokens=max_tokens,
        temperature=temperature,
        top_p=0.9,
        do_sample=True,
        repetition_penalty=1.1,
        streamer=streamer,
    )).start()

    partial = ""
    for token in streamer:
        partial += token
        yield partial

# ===== Gradio interface =====
demo = gr.ChatInterface(
    fn=respond,
    title="🤖 SAMAI",
    description="Cybersecurity & Programming AI — powered by Qwen2.5-14B",
    additional_inputs=[
        gr.Slider(64, 2048, value=512, step=64, label="Max tokens"),
        gr.Slider(0.1, 1.5, value=0.7, step=0.05, label="Temperature"),
    ],
    additional_inputs_accordion=gr.Accordion("⚙️ Settings", open=False),
)

demo.launch()