import gradio as gr
from time import sleep
from pymongo import MongoClient
from bson import ObjectId
from openai import OpenAI
openai_client = OpenAI()
import os

uri = os.environ.get('MONGODB_ATLAS_URI')
client = MongoClient(uri)
db_name = 'whatscooking'
collection_name = 'restaurants'
restaurants_collection = client[db_name][collection_name]
trips_collection = client[db_name]['smart_trips']



def get_restaurants(search, location, meters):

    newTrip = pre_aggregate_meters(location, meters)

    response = openai_client.embeddings.create(
            input=search,
            model="text-embedding-3-small",
            dimensions=256
        )

    restaurant_docs = list(trips_collection.aggregate([{
        "$vectorSearch": {
            "index" : "vector_index",
            "queryVector": response.data[0].embedding,
            "path" : "embedding",
            "numCandidates": 10,
            "limit": 3,
            "filter": {"searchTrip": newTrip}
        }},
        {"$project": {"_id" : 0, "embedding": 0}}]))

    
    chat_response = openai_client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content": "You are a helpful restaurant assistant."},
            { "role": "user", "content": f"Find me the 2 best restaurant and why based on {search} and  {restaurant_docs}. explain trades offs and why I should go to each one."}
        ]
        )

    trips_collection.delete_many({"searchTrip": newTrip})

    return chat_response.choices[0].message.content
    

def pre_aggregate_meters(location, meters):

    tripId = ObjectId()

    restaurants_collection.aggregate([
        {
            "$geoNear": {
            "near": location,
            "distanceField": "distance",
            "maxDistance": meters,
            "spherical": True,
            },
        },
        {
            "$addFields": {
                "searchTrip" : tripId,
                "date" : tripId.generation_time
            }
        },
        {
            "$merge": {
                "into": "smart_trips"
            }
        }   
    ]);

    sleep(10)

    return tripId


with gr.Blocks() as demo:
    gr.Markdown(
    """
    # MongoDB's Vector Restaurant planner 
    Start typing below to see the results
    """)
    gr.HTML(value='<iframe style="background: #FFFFFF;border: none;border-radius: 2px;box-shadow: 0 2px 10px 0 rgba(70, 76, 79, .2);" width="640" height="480" src="https://charts.mongodb.com/charts-paveldev-wiumf/embed/charts?id=65c24b0c-2215-4e6f-829c-f484dfd8a90c&maxDataAge=3600&theme=light&autoRefresh=true"></iframe>')
#    
    gr.Interface(
        get_restaurants,
        [
         
            gr.Textbox(placeholder="What type of dinner are you looking for?"),
         gr.Radio([("work",{
            "type": "Point",
            "coordinates": [
            -73.98527039999999,
            40.7589099
            ]
        }), ("home",{
            "type": "Point",
            "coordinates": [
            40.701975, -74.013686
            ]
        }), ("park", {
            "type": "Point",
            "coordinates": [40.720777, -74.000468
            ]
        })], label="Location", info="What location you need?"),
        gr.Slider(minimum=500, maximum=10000, randomize=False, step=5, label="Radius in meters")],
       gr.Textbox(label="MongoDB Vector Recommendations", placeholder="Results will be displayed here"),
        
    )
    #radio.change(location_searched, loc, out)
if __name__ == "__main__":
    demo.launch()