import streamlit as st
import requests
import json
from typing import List, Dict

# --------------------- API CONFIG ---------------------
API_KEY = "fw_3Zb9GtEVuqZWPgQ7rxBqcmNn"  # or set env FIREWORKS_API_KEY
MODEL_ID = "accounts/sentientfoundation/models/dobby-unhinged-llama-3-3-70b-new"
FW_URL = "https://api.fireworks.ai/inference/v1/chat/completions"

# --------------------- LLM CALL -----------------------

def ask_dobby(history: List[Dict], max_tokens: int = 150, temperature: float = 0.9) -> str:
    payload = {
        "model": MODEL_ID,
        "max_tokens": max_tokens,
        "temperature": temperature,
        "top_p": 1,
        "top_k": 40,
        "messages": history,
    }
    headers = {
        "Accept": "application/json",
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}",
    }
    r = requests.post(FW_URL, headers=headers, data=json.dumps(payload), timeout=60)
    r.raise_for_status()
    return r.json()["choices"][0]["message"]["content"].strip()

# --------------------- UI SETUP ----------------------
st.set_page_config("Rizzy", "💬")

st.image("logo.png") 

if "chat" not in st.session_state:
    st.session_state.chat: List[Dict] = []

SYSTEM_PROMPT = ("""
You are a dating friend, designed to be a tinder/hinge/bumble chat response generator, an outrageously charismatic friend who gives authentic, witty, flirty lines. 
It should be SHORT, easygoing, light, concise, simple, casual,chill, warm, vibey, carefree, cool. The key is not to try too hard.
NEVER be generic. Keep the text message you give the user SUPER SHORT and use text acronyms when applicable (like 'ur' instead of 'your', m2, omw, ik, etc.). Only give one answer. 2 sentences at max, ever.
For your answers: give the text message the user needs to paste to the dating app FIRST, and then maybe your super short reasoning. 
Always start with: [Say:"the txt message"], then maybe your SUPER SHORT description.
""")

# print existing chat
for msg in st.session_state.chat:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])

if not st.session_state.chat:
    welcome = "👋 Yo, I'm Rizzy. I'm here to help you chat for your dating life 🔥."
    st.session_state.chat.append({"role": "assistant", "content": welcome})
    st.chat_message("assistant").markdown(welcome)

user_input = st.chat_input("What's the situation or reply?")
if user_input is not None:
    # append user message
    st.session_state.chat.append({"role": "user", "content": user_input})
    with st.chat_message("user"):
        st.markdown(user_input)

    # build history (system + last 10 turns)
    history = [{"role": "system", "content": SYSTEM_PROMPT}] + st.session_state.chat[-10:]
    try:
        answer = ask_dobby(history)
    except Exception as e:
        answer = f"[Error contacting Dobby: {e}]"

    st.session_state.chat.append({"role": "assistant", "content": answer})
    with st.chat_message("assistant"):
        st.markdown(answer)