import os
import streamlit as st
import cv2
import numpy as np
from PIL import Image
from ultralytics import YOLO
from groq import Groq

# Load API key from Streamlit secrets
groq_api_key = st.secrets["GROQ_API_KEY"]

# Initialize Groq client
client = Groq(api_key=groq_api_key)

# Load YOLOv8 model (pretrained on COCO dataset)
model = YOLO("yolov8n.pt")

# Function to get AI-powered recommendations
def get_ai_recommendation(detected_objects):
    if not detected_objects:
        return "No objects detected. Please try again with a clearer image."
    
    prompt = f"Suggest optimal placement and design improvements for: {', '.join(detected_objects)} in a room."
    
    try:
        chat_completion = client.chat.completions.create(
            messages=[{"role": "user", "content": prompt}], 
            model="llama-3.3-70b-versatile"
        )
        return chat_completion.choices[0].message.content
    except Exception as e:
        return f"Error fetching AI recommendation: {e}"

# Streamlit UI
st.title("🏡 AI-Powered Home Setup Chatbot")
st.write("Upload an image of your room or use your camera for live capture.")

# Upload image option
uploaded_file = st.file_uploader("Upload an image of your room:", type=["jpg", "jpeg", "png"])

# Live camera capture
use_camera = st.checkbox("Use Live Camera")

if use_camera:
    st.write("Click 'Capture Image' to take a photo.")
    camera = cv2.VideoCapture(0)  # Open the webcam
    
    if st.button("Capture Image"):
        ret, frame = camera.read()
        if ret:
            image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            captured_image = Image.fromarray(image)
            st.image(captured_image, caption="Captured Image", use_column_width=True)
            
            # Run YOLOv8 for object detection
            results = model(image)
            detected_objects = [model.names[int(box.cls)] for box in results[0].boxes]
            
            # Display detected objects
            st.subheader("Detected Objects:")
            st.write(", ".join(set(detected_objects)))

            # Get AI-powered recommendations
            st.subheader("AI-Powered Placement Suggestions:")
            ai_recommendations = get_ai_recommendation(detected_objects)
            st.write(ai_recommendations)
        else:
            st.error("Failed to capture image. Please check your camera and try again.")

    camera.release()
    cv2.destroyAllWindows()

elif uploaded_file:
    image = Image.open(uploaded_file)
    image = np.array(image)
    
    # Run YOLOv8 for object detection
    results = model(image)
    detected_objects = [model.names[int(box.cls)] for box in results[0].boxes]
    
    # Display uploaded image and detected objects
    st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
    st.subheader("Detected Objects:")
    st.write(", ".join(set(detected_objects)))

    # Get AI-powered recommendations
    st.subheader("AI-Powered Placement Suggestions:")
    ai_recommendations = get_ai_recommendation(detected_objects)
    st.write(ai_recommendations)

# Text-based recommendations
user_input = st.text_input("Or enter an object name:")
if user_input:
    ai_recommendation = get_ai_recommendation([user_input.lower()])
    st.subheader("AI-Powered Suggestion:")
    st.write(ai_recommendation)