# Embedding and similarity using OpenAI Embeddings API
#
# https://platform.openai.com/docs/guides/embeddings
# Model: text-embedding-3-large (3072 dimensions, best quality)
#        text-embedding-3-small (1536 dimensions, faster/cheaper)
#        text-embedding-ada-002  (1536 dimensions, legacy)
#
# API compatible with emb4 (emb_gte-large_flask.py)
# Endpoints: GET/POST /embed, POST /embedding, POST /similarity
#
# Configuration via environment variables:
#   OPENAI_API_KEY  - required, your OpenAI API key
#   OPENAI_MODEL    - optional, default: text-embedding-3-large
#   OPENAI_BASE_URL - optional, default: https://api.openai.com/v1
#                     override to use Azure OpenAI or compatible APIs
#
import numpy as np
import os
import logging
import time

from openai import OpenAI
from flask import Flask, abort, jsonify, request

logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO,
)
logger = logging.getLogger(__name__)

# Configuration from environment variables
OPENAI_API_KEY   = os.environ.get("OPENAI_API_KEY", "")
OPENAI_MODEL     = os.environ.get("OPENAI_MODEL", "text-embedding-3-small")
OPENAI_BASE_URL  = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")
OPENAI_DIMENSIONS = int(os.environ.get("OPENAI_DIMENSIONS", "1024"))

if not OPENAI_API_KEY:
    logger.warning("OPENAI_API_KEY is not set — all requests will fail!")

client = OpenAI(api_key=OPENAI_API_KEY, base_url=OPENAI_BASE_URL)

logger.info("OpenAI embedding service started")
logger.info("Model      : " + OPENAI_MODEL)
logger.info("Base URL   : " + OPENAI_BASE_URL)
logger.info("Dimensions : " + str(OPENAI_DIMENSIONS))


def round_floats(o):
    if isinstance(o, float): return round(o, 6)
    if isinstance(o, dict): return {k: round_floats(v) for k, v in o.items()}
    if isinstance(o, (list, tuple)): return [round_floats(x) for x in o]
    return o

def cosine_similarity(a, b):
    a = np.array(a)
    b = np.array(b)
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))

def _get_embedding(text):
    """Call OpenAI API for a single text, return embedding as list of floats."""
    response = client.embeddings.create(
        model=OPENAI_MODEL,
        input=text,
        encoding_format="float",
        dimensions=OPENAI_DIMENSIONS
    )
    return response.data[0].embedding

def _get_embeddings(texts):
    """Call OpenAI API for a batch of texts, return list of embeddings."""
    response = client.embeddings.create(
        model=OPENAI_MODEL,
        input=texts,
        encoding_format="float",
        dimensions=OPENAI_DIMENSIONS
    )
    # Preserve input order (OpenAI returns results with index field)
    results = sorted(response.data, key=lambda x: x.index)
    return [item.embedding for item in results]

def embed_query(query):
    """Embed a search query (same as document embedding for OpenAI models)."""
    return round_floats(_get_embedding(query))

def embed(text):
    """Embed a single document text."""
    return round_floats(_get_embedding(text))

def embedding(texts):
    """Embed a list of texts."""
    return round_floats(_get_embeddings(texts))

def similarity(text, texts):
    """Compute cosine similarity between a query and a list of texts."""
    query_embedding = _get_embedding(text)
    doc_embeddings = _get_embeddings(texts)
    return [cosine_similarity(query_embedding, doc_emb) for doc_emb in doc_embeddings]


#====================================
#http://localhost:8085/embed?text=test

app = Flask(__name__)
app.config['CORS_HEADERS'] = 'Content-Type'
app.config['DEBUG'] = False

@app.route("/embed", methods=['GET'])
def get_embed():
    isQuery = request.args.get('isQuery')
    text = request.args.get('text')
    if not text or len(text) == 0:
        abort(500)
    else:
        start = time.time()
        if isQuery == 'true':
            result = embed_query(text)
        else:
            result = embed(text)
        end = time.time()
        took = int((end - start) * 1000)
        logger.info('Embed took %d (ms)' % took)
        return jsonify({'result': result, 'took': took})

@app.route("/embed", methods=['POST'])
def post_embed():
    data = request.get_json()
    isQuery = data.get('isQuery', False)
    if 'text' not in data or len(data['text']) == 0:
        abort(500)
    else:
        start = time.time()
        text = data['text']
        if isQuery:
            result = embed_query(text)
        else:
            result = embed(text)
        end = time.time()
        took = int((end - start) * 1000)
        logger.info('Embed took %d (ms)' % took)
        return jsonify({'result': result, 'took': took})

@app.route("/embedding", methods=['POST'])
def post_embedding():
    data = request.get_json()
    if 'texts' not in data or len(data['texts']) == 0:
        abort(500)
    else:
        start = time.time()
        texts = data['texts']
        results = embedding(texts)
        end = time.time()
        took = int((end - start) * 1000)
        logger.info('Embedding %d texts took %d (ms)' % (len(texts), took))
        return jsonify({'result': results, 'took': took})

@app.route("/similarity", methods=['POST'])
def post_similarity():
    data = request.get_json()
    if 'text' not in data or len(data['text']) == 0 or 'texts' not in data or len(data['texts']) == 0:
        abort(500)
    else:
        start = time.time()
        text = data['text']
        texts = data['texts']
        result = similarity(text, texts)
        end = time.time()
        took = int((end - start) * 1000)
        logger.info('Similarity for %d texts took %d (ms)' % (len(texts), took))
        return jsonify({'result': result, 'took': took})

if __name__ == "__main__":
    app.run(host='0.0.0.0')
#    app.run(host="0.0.0.0", port=80, debug=False)
