import streamlit as st
from pymongo import MongoClient
import os
from dotenv import load_dotenv
from helper.upload_file_to_s3 import upload_file
from helper.process_image import process_image_using_llm
from helper.create_embeddings import create_embedding
import time

# Load environment variables
load_dotenv()
AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
AWS_BUCKET_NAME = os.getenv("AWS_BUCKET_NAME")
MONGO_URI = os.getenv("MONGO_URI")
DB_NAME = os.getenv("DB_NAME")
COLLECTION_NAME = os.getenv("COLLECTION_NAME")
COLLECTION_NAME2=os.getenv("COMPANY_COLLECTION_NAME")

mongo_client = MongoClient(MONGO_URI)
db = mongo_client[DB_NAME]
collection = db[COLLECTION_NAME]
collection2=db[COLLECTION_NAME2]

def upload():
    if st.button("Back"):
        st.session_state.page = "upload_main"
        st.rerun()

    # File uploader (image files only)
    uploaded_image = st.file_uploader("Choose an image file to upload", type=["png", "jpg", "jpeg"],
                                      accept_multiple_files=False)

    # Fetch tags and categories from MongoDB
    tags_doc = collection2.find_one({"type": "tags"})
    categories_doc = collection2.find_one({"type": "categories"})

    tags_options = tags_doc["tags"] if tags_doc and "tags" in tags_doc else []
    categories_options = categories_doc["categories"] if categories_doc and "categories" in categories_doc else []

    # Multi-select dropdowns for tags and categories
    selected_tags = st.multiselect("Select Tags", options=tags_options)
    selected_categories = st.multiselect("Select Categories", options=categories_options)

    if uploaded_image and selected_tags and selected_categories:
        flag=False
        if st.button("Submit"):

            with st.spinner(text="Uploading and Processing Image"):
                # Upload file to S3
                metadata = upload_file(uploaded_image,"Image")
                if metadata:
                    object_url = metadata.get("object_url")
                    filename = metadata.get("name")

                    # Process image with LLM for description
                    llm_processed = process_image_using_llm(object_url)
                    if llm_processed:
                        # Create embedding with tags and categories in metadata
                        embedding_created = create_embedding(
                            object_url,
                            selected_tags,
                            selected_categories
                        )
                        if embedding_created:
                            # Save tags and categories to MongoDB document for the uploaded image
                            collection.update_one(
                                {"object_url": object_url},
                                {"$set": {
                                    "tags": selected_tags,
                                    "categories": selected_categories,
                                    "status":"processed"
                                }}
                            )
                            st.success("Image has been successfully uploaded and processed.")
                            flag=True
                        else:
                            st.error("Could not create embedding. Please try again.")
                            collection.update_one(
                                {"object_url": object_url},
                                {"$set": {

                                    "status": "failed"
                                }}
                            )
                    else:
                        st.error("Could not process the image description. Please try again.")
                        st.error("Could not create embedding. Please try again.")
                        collection.update_one(
                            {"object_url": object_url},
                            {"$set": {

                                "status": "failed"
                            }}
                        )
                else:
                    st.error("Could not upload the image. Please try again.")


            if flag:
                st.write("Redirecting to View Page to view all uploaded images")
                time.sleep(2)
                st.session_state.page = "view_image"
                st.rerun()