frontend/streamlit_app.py

import streamlit as st
import requests
import os
import uuid

# === CONFIG ===
API_URL = os.getenv("API_URL", "http://localhost:8000")
DEFAULT_HEADERS = {}

st.set_page_config(page_title="🧠 LLaMA Local AI Agent", layout="wide")

# === SESSION STATE INIT ===
if "token" not in st.session_state:
    st.session_state.token = ""
if "session_id" not in st.session_state:
    st.session_state.session_id = str(uuid.uuid4())

# === AUTH HELPERS ===
def auth_headers():
    headers = DEFAULT_HEADERS.copy()
    if st.session_state.token:
        headers["Authorization"] = f"Bearer {st.session_state.token}"
    return headers

def post(endpoint, data=None, files=None, json=None):
    try:
        res = requests.post(f"{API_URL}{endpoint}", data=data, files=files, json=json, headers=auth_headers())
        res.raise_for_status()
        return res.json()
    except Exception as e:
        st.error(f"❌ API POST Error: {e}")
        return None

def get(endpoint, params=None):
    try:
        res = requests.get(f"{API_URL}{endpoint}", params=params, headers=auth_headers())
        res.raise_for_status()
        return res.json()
    except Exception as e:
        st.error(f"❌ API GET Error: {e}")
        return None

# === SIDEBAR AUTH ===
st.sidebar.title("🔐 Authentication")

# === LOGIN ===
st.sidebar.subheader("Login")
username = st.sidebar.text_input("Username", key="login_username")
password = st.sidebar.text_input("Password", type="password", key="login_password")
if st.sidebar.button("Login"):
    login_response = post("/login", data={"username": username, "password": password})
    if login_response and "access_token" in login_response:
        st.session_state.token = login_response["access_token"]
        st.success("✅ Logged in!")
        st.rerun()
    else:
        st.error("❌ Login failed. Check credentials.")

# === REGISTER ===
st.sidebar.markdown("---")
st.sidebar.subheader("🆕 New here? Register")
new_username = st.sidebar.text_input("New Username", key="reg_username")
new_password = st.sidebar.text_input("New Password", type="password", key="reg_password")
confirm_password = st.sidebar.text_input("Confirm Password", type="password", key="reg_confirm")
role = st.sidebar.selectbox("Role", ["user", "admin"], key="reg_role")

if st.sidebar.button("Register"):
    if new_password != confirm_password:
        st.sidebar.error("❌ Passwords do not match.")
    elif not new_username or not new_password:
        st.sidebar.warning("Please fill all fields.")
    else:
        reg_payload = {
            "username": new_username,
            "password": new_password,
            "role": role
        }
        reg_result = post("/register", json=reg_payload)
        if reg_result and "access_token" in reg_result:
            st.session_state.token = reg_result["access_token"]
            st.sidebar.success("✅ Registered & Logged in!")
            st.rerun()
        else:
            st.sidebar.error("❌ Registration failed (User may exist).")

# === IF LOGGED IN ===
if st.session_state.token:
    tabs = st.tabs(["💬 Chat", "📁 File Q&A", "🧠 Tools", "🖼️ Image AI", "🎤 Voice", "🤖 Agent", "📧 Email"])

    # === 💬 Chat Tab ===
    with tabs[0]:
        st.header("💬 Chat with AI")
        user_input = st.text_input("Ask something...")
        if st.button("Send", key="chat"):
            if user_input:
                result = post("/chat", json={
                    "session_id": st.session_state.session_id,
                    "user_message": user_input
                })
                if result:
                    st.markdown(f"**AI:** {result.get('bot_response')}")

    # === 📁 File Upload and QA Tab ===
    with tabs[1]:
        st.header("📁 Upload & Ask")
        file = st.file_uploader("Upload PDF, TXT or image", type=["pdf", "txt", "jpg", "jpeg", "png"])
        if file:
            st.success(f"✅ Uploaded: {file.name}")
            if file.type.startswith("image/"):
                st.image(file)
                with st.spinner("Processing..."):
                    caption = post("/image-caption", files={"file": file})
                    ocr = post("/ocr", files={"file": file})
                if caption:
                    st.info(f"🖼️ Caption: {caption.get('caption')}")
                if ocr:
                    st.info(f"📝 OCR: {ocr.get('text')}")
            else:
                upload_res = post("/upload", files={"file": file})
                if upload_res:
                    question = st.text_input("Ask a question about the file")
                    if st.button("Ask"):
                        answer = get("/query_file", params={"filename": file.name, "question": question})
                        if answer:
                            st.success(f"💡 Answer: {answer.get('answer')}")

    # === 🧠 Tools Tab ===
    with tabs[2]:
        st.header("🛠️ AI Tools")
        tool_input = st.text_input("Enter input")
        tool = st.selectbox("Select Tool", ["calc", "time", "joke", "search", "browse", "exec_python"])
        if st.button("Run Tool"):
            result = post(f"/tool/{tool}", json={"prompt": tool_input})
            if result:
                st.markdown(f"🧠 Result: {result.get('result')}")

    # === 🖼️ Image AI Tab ===
    with tabs[3]:
        st.header("🖼️ Caption Any Image")
        img_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
        if img_file:
            st.image(img_file, width=400)
            if st.button("Generate Caption"):
                caption = post("/caption_image", files={"file": img_file})
                if caption:
                    st.success(f"📝 Caption: {caption.get('caption')}")

    # === 🎤 Voice Tab ===
    with tabs[4]:
        st.header("🎤 Voice Chat with AI")
        voice_file = st.file_uploader("Upload MP3 or WAV", type=["mp3", "wav"])
        if voice_file:
            st.audio(voice_file)
            transcribed = post("/transcribe", files={"file": voice_file})
            if transcribed:
                st.success(f"📝 Transcription: {transcribed.get('transcription')}")
                if st.button("Respond"):
                    response = post("/chat", json={"message": transcribed["transcription"]})
                    if response:
                        reply = response.get("response")
                        st.success(f"🗣️ AI says: {reply}")
                        tts = requests.post(f"{API_URL}/speak", params={"text": reply}, headers=auth_headers())
                        if tts.status_code == 200:
                            st.audio(tts.content)
                        else:
                            st.error("❌ TTS failed.")

    # === 🤖 Agent Tab ===
    with tabs[5]:
        st.header("🤖 AI Agent")
        task = st.text_input("Enter an agent task")
        if st.button("Run Agent"):
            result = post("/agent", json={"prompt": task})
            if result:
                st.markdown(f"🛠️ Result: {result.get('result')}")
        st.divider()
        if st.button("📜 Export Chat History"):
            history = get("/history/export")
            if history:
                st.text_area("History", history.get("text", ""), height=300)

    # === 📧 Email Generator Tab ===
    with tabs[6]:
        st.header("📧 Promo Email Builder")
        product = st.text_input("Product")
        recipient = st.text_input("Recipient Email")
        discount = st.slider("Discount (%)", 5, 80, 15)
        if st.button("Generate Email"):
            email = get("/generate_email", params={"to": recipient, "product": product, "discount": discount})
            if email:
                st.code(email.get("email"))
 
else:
    st.warning("🔐 Please login or register using the sidebar to access the AI features.")