from flask import Flask, render_template
from flask_socketio import SocketIO
import pandas as pd
import numpy as np
import google.generativeai as genai
# import KEYS
import os

app = Flask(__name__)
socketio = SocketIO(app)

# genai.configure(api_key=KEYS.api_key.GOOGLE_API_KEY)
genai.configure(api_key=os.environ.get("GEMINI_API_KEY"))

n_steps = 10
df = pd.DataFrame()
for i in range(0,n_steps):
    df = pd.concat([df,pd.read_feather(f"data_{i}.feather")])
df.reset_index(inplace=True,drop=True)

def get_results(top_n = 6,query = "men shirt"):
    query_embedding = genai.embed_content(model="models/text-embedding-004",
                                        content=query,
                                        task_type="retrieval_query")['embedding']
    scores = df['embedding'].apply(lambda x: np.dot(x,query_embedding))
    scores = scores.sort_values(ascending=False)[0:top_n]
    return df.loc[scores.index][['productDisplayName','link']].to_numpy()


@app.route('/')
def index():
    return render_template('index.html')

@socketio.on("search")
def get_products(query):
    data = []
    for x in get_results(query=query):
        data.append({'url':x[1],'name':x[0]})
    socketio.emit('data',data)

if __name__ == '__main__':
    socketio.run(app,port=7860,allow_unsafe_werkzeug=True,host='0.0.0.0')