import os
import yaml
import uuid
import openai
import pinecone
from string import Template
from typing import List, Dict
from neo4j import GraphDatabase
from flask import Flask, request, abort


TEXT_EMBEDDINGS_MODEL = 'text-embedding-ada-002'
PINECONE_API_KEY = os.environ["PINECONE_API_KEY"]
NEO4J_ENDPOINT = os.environ["NEO4J_ENDPOINT"]
NEO4J_PASSWORD = os.environ["NEO4J_PASSWORD"]
openai.api_key = os.environ["OPENAI_API_KEY"]

with open('/data/namespaces.yaml') as f:
    namespaces = yaml.load(f, Loader=yaml.CLoader)

neo4j_driver = GraphDatabase.driver(NEO4J_ENDPOINT, auth=('neo4j', NEO4J_PASSWORD))
pc = pinecone.Pinecone(api_key=PINECONE_API_KEY)
pinecone_index = pc.Index(host="https://ossinsight-5c794ec.svc.aped-4627-b74a.pinecone.io")

app = Flask(__name__)

def get_embedding(text: str, model: str) -> List[float]:
    result = openai.Embedding.create(
        model=model,
        input=text,
    )
    return result["data"][0]["embedding"]

# the function used to get embedding
def get_text_embedding(text: str) -> List[float]:
    return get_embedding(text, TEXT_EMBEDDINGS_MODEL)

@app.route("/ping")
def ping():
    return {'status': 'success'}

@app.route("/<namespace>/query", methods=['POST'])
def query(namespace):
    if namespace not in namespaces:
        return {'status': 'error', 'message': f'no such namespace: {namespace}'}, 404
        
    query = request.json['query']
    filter = request.json.get('filter')
    top_k = request.json.get('top_k', 5)
    include_values = request.json.get('include_values', False)

    query_embedding = get_text_embedding(query)

    matches = [{
        **x['metadata'], 
        '_id': x['id'], 
        '_score': x['score'], 
        '_values': x.get('values'),
    } for x in pinecone_index.query(
        namespace=namespace,
        top_k=top_k,
        filter=filter,
        include_metadata=True,
        vector=query_embedding,
        include_values=include_values,
    )['matches']]
    
    return {'status': 'success', 'matches': matches}

@app.route("/<namespace>/upsert", methods=['PUT'])
def upsert(namespace):
    if namespace not in namespaces:
        return {'status': 'error', 'message': f'no such namespace: {namespace}'}, 404
    
    ns = namespaces[namespace]
    data = request.json
    
    for k in ns['index_keys'] + [ns['embedding_key']]:
        if k not in data:
            return {'status': 'error', 'message': f'key missing in data: {k}'}, 400
    
    index = ', '.join([f'{k}: ${k}' for k in ns['index_keys']])
    fields = ', '.join([f'x.{k} = ${k}' for k in data if k not in ns['index_keys']])
        
    clsname = ''.join([x.title() for x in namespace.split('_')])
    template = Template('merge (x:$clsname { $index }) on match set $fields on create set $fields return id(x) as id')

    query = template.substitute({
        'clsname': clsname,
        'index': index,
        'fields': fields,
    })
    
    with neo4j_driver.session(database="neo4j") as session:
        rs = session.run(query, **data)
        _id = str(uuid.UUID(fields=(0, 0, 0, 0, 0, [r.get('id') for r in rs][0])))
    
    pinecone_index.upsert(
        vectors=[
            {
                'id': _id,
                'values': get_text_embedding(data[ns['embedding_key']]),
                'metadata': data,
            }
        ],
        namespace=namespace,
    )
    
    return {'status': 'success', 'id': _id}

@app.route("/<namespace>/<vec_id>", methods=['DELETE'])
def delete(namespace, vec_id):
    if namespace not in namespaces:
        return {'status': 'error', 'message': f'no such namespace: {namespace}'}, 404
    
    node_id = uuid.UUID(vec_id).node
    clsname = ''.join([x.title() for x in namespace.split('_')])
    query = f'match (x:{clsname}) where id(x) = $id detach delete x'
    print(query)
    
    with neo4j_driver.session(database="neo4j") as session:
        session.run(query, id=node_id)
        
    pinecone_index.delete(ids=[vec_id], namespace=namespace)
    
    return {'status': 'success'}