app.py

import streamlit as st
from openai import OpenAI
from PIL import Image
import time
import os
import uuid
import firebase_admin
from firebase_admin import credentials, firestore
import re

# 🔐 Firebase setup
if not firebase_admin._apps:
    cred = credentials.Certificate("firebase-service-account.json")
    firebase_admin.initialize_app(cred)

db = firestore.client()

# 🔐 OpenAI setup
openai_key = os.getenv("openai_key")
assistant_id = os.getenv("ASSISTANT_ID")
client = OpenAI(api_key=openai_key)

# 🌐 Streamlit Config
st.set_page_config(page_title="Carfind.co.za AI Assistant", layout="wide")

# 🎯 Session + User ID
if "user_id" not in st.session_state:
    st.session_state["user_id"] = str(uuid.uuid4())
user_id = st.session_state["user_id"]

# 🖼️ Branding + Styling
st.markdown("""
    <style>
    .block-container {padding-top: 1rem; padding-bottom: 0rem;}
    header {visibility: hidden;}
    .stChatMessage { max-width: 85%; border-radius: 12px; padding: 8px; margin-bottom: 10px; }
    .stChatMessage[data-testid="stChatMessage-user"] { background: #f0f0f0; color: #000000; }
    .stChatMessage[data-testid="stChatMessage-assistant"] { background: #D6E9FE; color: #000000; }
    @keyframes bounceIn {
        0% { transform: scale(0.7); opacity: 0; }
        60% { transform: scale(1.1); opacity: 1; }
        80% { transform: scale(0.95); }
        100% { transform: scale(1); }
    }
    .carfind-logo {
        animation: bounceIn 0.6s ease-out;
    }
    .spec-table {
        width: 100%;
        border-collapse: collapse;
        margin-top: 20px;
    }
    .spec-table th {
        background-color: #0071BC;
        color: white;
        padding: 10px;
        text-align: left;
    }
    .spec-table td {
        padding: 10px;
        border: 1px solid #ccc;
        background-color: #f9f9f9;
    }
    .icon {
        width: 16px;
        vertical-align: middle;
        margin-right: 6px;
    }
    </style>
""", unsafe_allow_html=True)

st.markdown("""
    <div style='text-align: center; margin-top: 20px; margin-bottom: -10px;'>
        <span style='display: inline-flex; align-items: center; gap: 8px;'>
            <img src='https://www.carfind.co.za/images/Carfind-Icon.svg' width='30' class='carfind-logo'/>
            <span style='font-size: 14px; color: gray;'>Powered by Carfind</span>
        </span>
    </div>
""", unsafe_allow_html=True)

# 🔁 Firebase Chat Functions
def get_or_create_thread_id():
    doc_ref = db.collection("users").document(user_id)
    doc = doc_ref.get()
    if doc.exists:
        return doc.to_dict()["thread_id"]
    else:
        thread = client.beta.threads.create()
        doc_ref.set({
            "thread_id": thread.id,
            "created_at": firestore.SERVER_TIMESTAMP
        })
        return thread.id

def save_message(role, content):
    db.collection("users").document(user_id).collection("messages").add({
        "role": role,
        "content": content,
        "timestamp": firestore.SERVER_TIMESTAMP
    })

def display_chat_history():
    messages = db.collection("users").document(user_id).collection("messages") \
        .order_by("timestamp").stream()

    assistant_icon_html = "<img src='https://www.carfind.co.za/images/Carfind-Icon.svg' width='22' style='vertical-align:middle;'/>"
    for msg in list(messages)[::-1]:
        data = msg.to_dict()
        if data["role"] == "user":
            st.markdown(
                f"<div class='stChatMessage' data-testid='stChatMessage-user'>"
                f"👤 <strong>You:</strong> {data['content']}</div>", unsafe_allow_html=True
            )
        else:
            st.markdown(
                f"<div class='stChatMessage' data-testid='stChatMessage-assistant'>"
                f"{assistant_icon_html} <strong>Carfind Assistant:</strong> {data['content']}</div>",
                unsafe_allow_html=True
            )

# 🛍 Tabs: AI Chat | Car Identifier
tab1, tab2 = st.tabs(["AI Chat", "What car is that?"])

# 🤖 AI Chat Tab
with tab1:
    input_col, clear_col = st.columns([9, 1])
    with input_col:
        user_input = st.chat_input("Type your message here...")

    with clear_col:
        if st.button("🗑️", key="clear-chat", help="Clear Chat"):
            try:
                user_doc_ref = db.collection("users").document(user_id)
                for msg in user_doc_ref.collection("messages").stream():
                    msg.reference.delete()
                user_doc_ref.delete()
                st.session_state.clear()
                st.rerun()
            except Exception as e:
                st.error(f"Failed to clear chat: {e}")

    thread_id = get_or_create_thread_id()
    display_chat_history()

    if user_input:
        client.beta.threads.messages.create(thread_id=thread_id, role="user", content=user_input)
        save_message("user", user_input)

        with st.spinner("Thinking and typing... 💭"):
            run = client.beta.threads.runs.create(thread_id=thread_id, assistant_id=assistant_id)
            while True:
                run_status = client.beta.threads.runs.retrieve(thread_id=thread_id, run_id=run.id)
                if run_status.status == "completed":
                    break
                time.sleep(1)

        messages_response = client.beta.threads.messages.list(thread_id=thread_id)
        latest_response = sorted(messages_response.data, key=lambda x: x.created_at)[-1]
        assistant_message = latest_response.content[0].text.value
        save_message("assistant", assistant_message)
        time.sleep(0.5)
        st.rerun()

# 🚗 "What car is that?" Tab
with tab2:
    uploaded_image = st.file_uploader("Upload an image of a car and let Ai identify it for you", type=["jpg", "jpeg", "png"])

    if uploaded_image:
        col1, col2 = st.columns([1.2, 1.8])

        with col1:
            image = Image.open(uploaded_image)
            st.image(image, caption="Uploaded Image", use_container_width=True)

        with col2:
            try:
                image_thread = client.beta.threads.create()
                file_response = client.files.create(file=uploaded_image, purpose="assistants")

                client.beta.threads.messages.create(
                    thread_id=image_thread.id,
                    role="user",
                    content=[
                        {
                            "type": "image_file",
                            "image_file": {"file_id": file_response.id}
                        },
                        {
                            "type": "text",
                            "text": "Please identify this car from the image and include a full vehicle specification table, a short overview of the car, and who it's recommended for."
                        }
                    ]
                )

                run = client.beta.threads.runs.create(
                    thread_id=image_thread.id,
                    assistant_id=assistant_id
                )

                with st.spinner("🔍 Analyzing image and identifying the car..."):
                    while True:
                        run_status = client.beta.threads.runs.retrieve(
                            thread_id=image_thread.id,
                            run_id=run.id
                        )
                        if run_status.status == "completed":
                            break
                        time.sleep(1)

                messages = client.beta.threads.messages.list(thread_id=image_thread.id)
                assistant_message = messages.data[0].content[0].text.value

                cleaned_message = re.sub(r'\*\*(.*?)\*\*', r'<strong>\1</strong>', assistant_message)
                cleaned_message = cleaned_message.replace("###", "<h4 style='margin-bottom: 10px;'>").replace("\n", "<br>")

                # Replace key labels with icons
                cleaned_message = cleaned_message.replace("Transmission", "<img src='https://cdn-icons-png.flaticon.com/128/3448/3448449.png' class='icon'/> Transmission")
                cleaned_message = cleaned_message.replace("Engine Size", "<img src='https://cdn-icons-png.flaticon.com/128/888/888879.png' class='icon'/> Engine Size")

                st.success("✅ Identification Complete")
                st.markdown(f"""
                    <div style='
                        font-family: "Segoe UI", sans-serif;
                        background-color: #ffffff;
                        color: #003B6F;
                        padding: 24px;
                        border-radius: 12px;
                        box-shadow: 0 4px 12px rgba(0,0,0,0.05);
                        border-left: 5px solid #0071BC;
                        font-size: 15.5px;
                        line-height: 1.7;
                    '>
                        {cleaned_message}
                    </div>
                """, unsafe_allow_html=True)

            except Exception as e:
                st.error(f"❌ Error during image analysis: {str(e)}")