import streamlit as st
import cohere
import os
import base64

st.set_page_config(page_title="Cohere Chat", layout="wide")

AI_PFP = "media/pfps/cohere-pfp.png"
USER_PFP = "media/pfps/user-pfp.jpg"
BANNER = "media/banner.png"

if not os.path.exists(AI_PFP) or not os.path.exists(USER_PFP):
    st.error("Missing profile pictures in media/pfps directory")
    st.stop()

model_info = {
    "c4ai-aya-expanse-8b": {"description": "Aya Expanse is a highly performant 8B multilingual model...", "context": "4K", "output": "4K"},
    "c4ai-aya-expanse-32b": {"description": "Aya Expanse is a highly performant 32B multilingual model...", "context": "128K", "output": "4K"},
    "command-a-03-2025": {"description": "Command A is our most performant model to date...", "context": "256K", "output": "8K"},
    "command-r7b-12-2024": {"description": "command-r7b-12-2024 is a small, fast update...", "context": "128K", "output": "4K"},
    "command-r-plus-04-2024": {"description": "Command R+ is an instruction-following conversational model...", "context": "128K", "output": "4K"},
    "c4ai-aya-vision-8b": {"description": "Aya Vision is an 8B vision-language model...", "context": "4K", "output": "4K"},
    "c4ai-aya-vision-32b": {"description": "Aya Vision is a 32B vision-language model...", "context": "128K", "output": "4K"}
}

with st.sidebar:
    st.image(BANNER, use_container_width=True)
    st.markdown("Hugging Face 🤗 Community UI (Vision Model support coming soon)")
    st.title("Settings")
    api_key = st.text_input("Cohere API Key", type="password")
    selected_model = st.selectbox("Model", options=list(model_info.keys()))
    def clear_chat():
        st.session_state.messages = []
        st.session_state.first_message_sent = False
    st.button("Clear Chat", on_click=clear_chat)
    st.divider()
    st.image(AI_PFP, width=60)
    st.subheader(selected_model)
    st.markdown(model_info[selected_model]["description"])
    st.caption(f"Context: {model_info[selected_model]['context']}")
    st.caption(f"Output: {model_info[selected_model]['output']}")
    st.markdown("Powered by Cohere's API")

if "messages" not in st.session_state:
    st.session_state.messages = []
if "first_message_sent" not in st.session_state:
    st.session_state.first_message_sent = False
if "uploaded_image" not in st.session_state:
    st.session_state.uploaded_image = None

if not st.session_state.first_message_sent:
    st.markdown(
        "<h1 style='text-align:center; color:#4a4a4a; margin-top:100px;'>How can Cohere help you today?</h1>",
        unsafe_allow_html=True
    )

for msg in st.session_state.messages:
    avatar = USER_PFP if msg["role"] == "user" else AI_PFP
    with st.chat_message(msg["role"], avatar=avatar):
        st.markdown(msg["content"])

col1, col2 = st.columns([1, 8])
with col1:
    if selected_model.startswith("c4ai-aya-vision"):
        img = st.file_uploader(label="📷", key="uploader", type=["png","jpg","jpeg"], accept_multiple_files=False)
        if img is not None:
            st.session_state.uploaded_image = img
            st.image(img, width=80)
    else:
        st.write("")
with col2:
    prompt = st.chat_input("Message...")

if prompt or st.session_state.uploaded_image:
    if not api_key:
        st.error("API key required")
        st.stop()
    user_items = []
    if prompt:
        st.session_state.first_message_sent = True
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user", avatar=USER_PFP):
            st.markdown(prompt)
        user_items.append({"type": "text", "text": prompt})
    if st.session_state.uploaded_image:
        raw = st.session_state.uploaded_image.read()
        b64 = base64.b64encode(raw).decode("utf-8")
        url = f"data:image/jpeg;base64,{b64}"
        user_items.append({"type": "image_url", "image_url": {"url": url}})
        with st.chat_message("user", avatar=USER_PFP):
            st.image(raw, width=200)
        st.session_state.uploaded_image = None
    try:
        co = cohere.ClientV2(api_key)
        response = co.chat(model=selected_model, messages=[{"role":"user","content":user_items}])
        reply = "".join(getattr(item,'text','') for item in response.message.content)
        with st.chat_message("assistant", avatar=AI_PFP):
            st.markdown(reply)
        st.session_state.messages.append({"role":"assistant","content":reply})
    except Exception as e:
        st.error(f"Error: {e}")