import streamlit as st
from pymongo import MongoClient
import os
from pydantic import BaseModel
from typing import List


# Load environment variables
MONGO_URI = os.getenv("MONGO_URI")
MONGO_DB_NAME = os.getenv("MONGO_DB_NAME")

# Connect to MongoDB
client = MongoClient(MONGO_URI)
#db = client[MONGO_DB_NAME]
db = client["vector_db"]
#vectors_collection = db["vector_collection"]
vectors_collection = db["vector_collection"]

 #client = MongoClient(connection_string, tls=True, tlsAllowInvalidCertificates=True)
 #   db = client[db_name]
 #   collection = db[collection_name]

# Define a Pydantic model for vectors
class Vector(BaseModel):
    vector: List[float]
    metadata: dict

# Streamlit app
st.title("MongoDB Vector Search with Streamlit")

# Form to insert a vector
with st.form("vector_form"):
    st.write("Insert a new vector")
    vector_input = st.text_input("Enter vector (comma-separated floats):")
    metadata_input = st.text_input("Enter metadata (JSON format):")
    submitted = st.form_submit_button("Submit")

    if submitted:
        try:
            # Parse the vector and metadata
            vector = [float(x) for x in vector_input.split(",")]
            metadata = eval(metadata_input)  # Convert JSON string to dict

            # Create a Vector object
            new_vector = Vector(vector=vector, metadata=metadata)

            # Insert into MongoDB
            result = vectors_collection.insert_one(new_vector.dict())
            st.success(f"Vector inserted with ID: {result.inserted_id}")
        except Exception as e:
            st.error(f"Error: {e}")

# Display all vectors
st.write("## Stored Vectors")
vectors = vectors_collection.find()
for vector in vectors:
    st.json(vector)