import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import torch
from collections import defaultdict
from datetime import datetime
import pandas as pd

# PDF export
from reportlab.pdfgen import canvas
from reportlab.lib.pagesizes import letter
import tempfile

# Model & sentiment
MODEL_NAME = "HuggingFaceTB/SmolLM2-360M-Instruct"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)  # NO device_map here
sentiment_analyzer = pipeline("sentiment-analysis")

SYSTEM_PROMPT = "You are a friendly assistant with fire vibes."

feedback_store = []
like_leaderboard = defaultdict(int)

def now_str():
    return datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S UTC")

def generate_reply(prompt_str, temp, max_tokens):
    inputs = tokenizer(prompt_str, return_tensors="pt")
    outputs = model.generate(
        **inputs,
        max_new_tokens=max_tokens,
        temperature=temp,
        do_sample=True
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

def respond(message, history, user, temp, max_tokens):
    username = user.name if user else "anonymous"
    prompt = SYSTEM_PROMPT + "\n"
    for msg_info in history:
        prompt += f"[{msg_info['timestamp']}] {msg_info['user']}: {msg_info['user_msg']}\n"
        prompt += f"[{msg_info['timestamp']}] 🧠 {msg_info['bot_msg']}\n"
    prompt += f"[{now_str()}] {username}: {message}\nAssistant:"

    bot_reply = generate_reply(prompt, temp, max_tokens)
    sentiment = sentiment_analyzer(bot_reply)[0]
    icon = "😊" if sentiment["label"] == "POSITIVE" else "😐" if sentiment["label"] == "NEUTRAL" else "☹️"
    record = {
        "timestamp": now_str(),
        "user": username,
        "user_msg": message,
        "bot_msg": bot_reply,
        "sentiment": sentiment,
        "icon": icon
    }
    history.append(record)
    return history

def record_feedback(history, fb):
    if history:
        last = history[-1]
        feedback_store.append({
            "timestamp": last["timestamp"],
            "user_msg": last["user_msg"],
            "bot_msg": last["bot_msg"],
            "sentiment": last["sentiment"],
            "feedback": fb
        })
        if fb == "like":
            like_leaderboard[last["bot_msg"]] += 1
    return history

def download_csv():
    df = pd.DataFrame(feedback_store)
    path = "/tmp/sentiment_feedback.csv"
    df.to_csv(path, index=False)
    return path

def leaderboard_text():
    sorted_leader = sorted(like_leaderboard.items(), key=lambda x: x[1], reverse=True)
    lines = [f"{i+1}. {msg[:60]}... → {count} likes" for i,(msg,count) in enumerate(sorted_leader)]
    return "\n".join(lines)

def export_pdf(history):
    temp_pdf = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf").name
    c = canvas.Canvas(temp_pdf, pagesize=letter)
    textobj = c.beginText(40, 750)
    textobj.setFont("Helvetica", 10)

    for msg in history:
        textobj.textLine(f"[{msg['timestamp']}] {msg['user']}: {msg['user_msg']}")
        textobj.textLine(f"   🧠 {msg['icon']}: {msg['bot_msg']}")
        textobj.textLine("")

    c.drawText(textobj)
    c.save()
    return temp_pdf

with gr.Blocks() as demo:
    gr.Markdown("## 🔥 Smol Chatbot 🔥")

    login_button = gr.LoginButton()
    history_state = gr.State([])

    with gr.Row():
        temp_slider = gr.Slider(0.1, 1.2, value=0.7, label="Temperature")
        max_tokens_slider = gr.Slider(20, 300, value=150, label="Max Tokens")

    msg = gr.Textbox(label="Your message")
    send = gr.Button("Send")
    like = gr.Button("👍 Like")
    dislike = gr.Button("👎 Dislike")
    download_csv_btn = gr.Button("📥 Download CSV")
    download_pdf_btn = gr.Button("📄 Download PDF")
    leaderboard_btn = gr.Button("🏆 Leaderboard")
    leaderboard_out = gr.Textbox(label="Leaderboard")

    def on_send(message, history, user, temp_val, max_toks):
        new_hist = respond(message, history, user, temp_val, max_toks)
        return "", new_hist

    send.click(
        on_send,
        [msg, history_state, login_button, temp_slider, max_tokens_slider],
        [msg, history_state],
    )

    like.click(record_feedback, history_state, history_state, _js="() => ['like']")
    dislike.click(record_feedback, history_state, history_state, _js="() => ['dislike']")

    download_csv_btn.click(download_csv, None, gr.File())
    download_pdf_btn.click(export_pdf, history_state, gr.File())
    leaderboard_btn.click(lambda: leaderboard_text(), None, leaderboard_out)

demo.launch()