import os
import uuid
from string import Template
from typing import Dict, List

import openai
import pinecone
import yaml
from flask import Flask, abort, request
from neo4j import GraphDatabase

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))
pinecone.init(api_key=PINECONE_API_KEY, environment="us-east1-gcp")
pinecone_index = pinecone.Index("ossinsight-index")

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)

    query_embedding = get_text_embedding(query)

    matches = [
        {**x["metadata"], "_id": x["id"], "_score": x["score"]}
        for x in pinecone_index.query(
            namespace=namespace,
            top_k=top_k,
            filter=filter,
            include_metadata=True,
            vector=query_embedding,
        )["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"

    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"}