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

MODEL_ID = "Jeppcode/ScalableLab2"
SUBFOLDER = "merged-model-fp16"

print(f"Loading model {MODEL_ID}/{SUBFOLDER} ...")

# Tokenizer
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID, subfolder=SUBFOLDER,
)

# Model – fp16 and optimized for CPU
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, 
    subfolder=SUBFOLDER, 
    dtype=torch.float16, 
    low_cpu_mem_usage=True, 
    device_map="cpu",
)
model.eval()

# Hardcoded system prompt
SYSTEM_PROMPT = "You are a helpful, polite assistant. Give clear and structured explanations."

def build_prompt(message, history):
    messages = []
    messages.append({"role": "system", "content": SYSTEM_PROMPT})
    
    for msg in history:
        role = msg.get("role")
        content = msg.get("content", "")
        if isinstance(content, list):
            texts = []
            for block in content:
                if isinstance(block, dict) and block.get("type") == "text":
                    texts.append(block.get("text", ""))
                else:
                    texts.append(str(block))
            text = "\n".join(t for t in texts if t)
        else:
            text = str(content)
            
        if text and role in ("user", "assistant", "system"):
            messages.append({"role": role, "content": text})
            
    messages.append({"role": "user", "content": message})
    
    prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True,
    )
    return prompt

def chat_fn(message, history, max_new_tokens):
    prompt = build_prompt(message, history)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    # Default values for removed sliders
    temperature = 0.7 
    top_p = 0.9
    repetition_penalty = 1.1 

    gen_kwargs = {
        **inputs,
        "max_new_tokens": int(max_new_tokens),
        "pad_token_id": tokenizer.eos_token_id,
        "eos_token_id": tokenizer.eos_token_id,
        "repetition_penalty": float(repetition_penalty),
        "do_sample": True,
        "temperature": float(temperature),
        "top_p": float(top_p),
    }

    with torch.no_grad():
        outputs = model.generate(**gen_kwargs)
        generated = tokenizer.decode(
            outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True,
        ).strip()

    return generated

# --- Nature Theme Configuration ---
nature_theme = gr.themes.Soft(
    primary_hue="green",
    secondary_hue="emerald",
    neutral_hue="stone",
).set(
    body_background_fill="#f5f7f5",
    block_background_fill="rgba(255, 255, 255, 0.9)",
    button_primary_background_fill="#2E7D32", 
    border_color_primary="#4CAF50",
)

custom_css = """
.gradio-container {
    background: url('https://images.unsplash.com/photo-1441974231531-c6227db76b6e?q=80&w=2560&auto=format&fit=crop') no-repeat center center fixed; 
    background-size: cover;
}
footer {visibility: hidden}
"""

max_new_tokens_slider = gr.Slider(
    minimum=16, maximum=512, value=128, step=8, label="Response Length (Tokens)",
)

# Instantiate ChatInterface WITHOUT theme/css arguments initially to avoid TypeError
demo = gr.ChatInterface(
    fn=chat_fn,
    title="🌿 NatureChat Lab 2",
    description="Chat with the fine-tuned Llama model. Relax and enjoy the view.",
    additional_inputs=[max_new_tokens_slider],
    additional_inputs_accordion="Settings"
)

# Apply theme and CSS manually after instantiation
demo.theme = nature_theme
demo.css = custom_css

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