# 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 — dimensions param not supported)
#
# 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_EMB_MODEL            - optional, default: text-embedding-3-small
#   OPENAI_BASE_URL             - optional, default: https://api.openai.com/v1 override to use Azure OpenAI or compatible APIs
#   OPENAI_EMB_MODEL_DIMENSIONS - optional, default: 1024 (ignored for ada-002)

import logging
import os
import time

import numpy as np
import openai
from flask import Flask, abort, jsonify, request
from openai import OpenAI

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

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

OPENAI_API_KEY    = os.environ.get("OPENAI_API_KEY", "")
OPENAI_MODEL      = os.environ.get("OPENAI_EMB_MODEL", "text-embedding-3-small")
OPENAI_BASE_URL   = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")
OPENAI_MODEL_DIMENSIONS = int(os.environ.get("OPENAI_EMB_MODEL_DIMENSIONS", "1024"))

# dimensions parameter is only supported by text-embedding-3-* models
SUPPORTS_DIMENSIONS = OPENAI_MODEL.startswith("text-embedding-3-")

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      : %s", OPENAI_MODEL)
logger.info("Base URL   : %s", OPENAI_BASE_URL)
logger.info("Dimensions : %s", OPENAI_MODEL_DIMENSIONS if SUPPORTS_DIMENSIONS else "default (model does not support override)")

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def is_query(val) -> bool:
    if isinstance(val, bool):
        return val
    return str(val).lower() == 'true'


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, b = np.array(a), np.array(b)
    norm = np.linalg.norm(a) * np.linalg.norm(b)
    if norm == 0:
        return 0.0
    return float(np.dot(a, b) / norm)


# ---------------------------------------------------------------------------
# OpenAI API calls
# ---------------------------------------------------------------------------

def get_embeddings(texts: list) -> list:
    """Batch embed a list of texts. Returns embeddings in input order."""
    kwargs = dict(model=OPENAI_MODEL, input=texts, encoding_format="float")
    if SUPPORTS_DIMENSIONS:
        kwargs["dimensions"] = OPENAI_MODEL_DIMENSIONS
    response = client.embeddings.create(**kwargs)
    results = sorted(response.data, key=lambda x: x.index)
    return [item.embedding for item in results]


# ---------------------------------------------------------------------------
# Encoding functions
# ---------------------------------------------------------------------------

def embed_query(query: str):
    return round_floats(get_embeddings([query])[0])


def embed(text: str):
    return round_floats(get_embeddings([text])[0])


def embedding(texts: list):
    return round_floats(get_embeddings(texts))


def similarity(text: str, texts: list):
    # single batched API call for both query and docs
    all_embeddings = get_embeddings([text] + texts)
    query_emb = all_embeddings[0]
    doc_embs = all_embeddings[1:]
    return [cosine_similarity(query_emb, doc_emb) for doc_emb in doc_embs]


# ---------------------------------------------------------------------------
# Flask app
# ---------------------------------------------------------------------------

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

app = Flask(__name__)
app.config['DEBUG'] = False


@app.route("/embed", methods=['GET'])
def get_embed():
    text = request.args.get('text') or ''
    if not text:
        abort(400, description="Missing or empty 'text' query parameter")
    try:
        start = time.perf_counter()
        result = embed_query(text) if is_query(request.args.get('isQuery')) else embed(text)
        took = int((time.perf_counter() - start) * 1000)
        logger.info('Embed took %d ms' % took)
        return jsonify({'result': result, 'took': took})
    except openai.OpenAIError as e:
        logger.error("OpenAI API error: %s", e)
        abort(502, description=str(e))


@app.route("/embed", methods=['POST'])
def post_embed():
    data = request.get_json()
    if not data or not data.get('text'):
        abort(400, description="Request must include a non-empty 'text' field")
    try:
        start = time.perf_counter()
        text = data['text']
        result = embed_query(text) if is_query(data.get('isQuery', False)) else embed(text)
        took = int((time.perf_counter() - start) * 1000)
        logger.info('Embed took %d ms' % took)
        return jsonify({'result': result, 'took': took})
    except openai.OpenAIError as e:
        logger.error("OpenAI API error: %s", e)
        abort(502, description=str(e))


@app.route("/embedding", methods=['POST'])
def post_embedding():
    data = request.get_json()
    if not data or not data.get('texts'):
        abort(400, description="Request must include a non-empty 'texts' field")
    try:
        texts = data['texts']
        start = time.perf_counter()
        results = embedding(texts)
        took = int((time.perf_counter() - start) * 1000)
        logger.info('Embedding %d texts took %d ms' % (len(texts), took))
        return jsonify({'result': results, 'took': took})
    except openai.OpenAIError as e:
        logger.error("OpenAI API error: %s", e)
        abort(502, description=str(e))


@app.route("/similarity", methods=['POST'])
def post_similarity():
    data = request.get_json()
    if not data or not data.get('text') or not data.get('texts'):
        abort(400, description="Request must include non-empty 'text' and 'texts' fields")
    try:
        start = time.perf_counter()
        result = similarity(data['text'], data['texts'])
        took = int((time.perf_counter() - start) * 1000)
        logger.info('Similarity for %d texts took %d ms' % (len(data['texts']), took))
        return jsonify({'result': result, 'took': took})
    except openai.OpenAIError as e:
        logger.error("OpenAI API error: %s", e)
        abort(502, description=str(e))


@app.route("/health", methods=['GET'])
def health():
    return jsonify({"status": "ok", "model": OPENAI_MODEL, "base_url": OPENAI_BASE_URL})


if __name__ == "__main__":
    app.run(host='0.0.0.0')
