# Embedding and similarity using Qwen3-Embedding-0.6B model
#
# https://huggingface.co/Qwen/Qwen3-Embedding-0.6B
# Requires transformers>=4.51.0, sentence-transformers>=2.7.0
#
# Instruction-aware: queries use prompt_name="query" (adds task instruction prefix)
# Documents are encoded as plain text (no instruction needed)
# Native 1024 dimensions, 32k context length

import logging
import os
import threading
import time

import numpy as np
import torch
from flask import Flask, abort, jsonify, request
from sentence_transformers import SentenceTransformer
from transformers import set_seed

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__)

seed_value = 1
set_seed(seed_value)
torch.manual_seed(seed_value)
if torch.cuda.is_available():
    torch.cuda.manual_seed(seed_value)
    torch.cuda.manual_seed_all(seed_value)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    logger.info("Cuda available!")

# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------

MODEL_NAME = os.environ.get("EMB_MODEL", "Qwen/Qwen3-Embedding-0.6B")

logger.info("Loading model: %s", MODEL_NAME)
model = SentenceTransformer(
    MODEL_NAME,
    model_kwargs={
        "device_map": "auto",
        "torch_dtype": torch.bfloat16,
    },
    tokenizer_kwargs={"padding_side": "left"},
)
logger.info("Model loaded: %s", MODEL_NAME)

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

cache_lock = threading.Lock()
cache_cnt = 0

def empty_cache():
    global cache_cnt
    with cache_lock:
        cache_cnt += 1
        flush = (cache_cnt % 10 == 0)
    if flush and torch.cuda.is_available():
        torch.cuda.empty_cache()


def is_query(val) -> bool:
    """Accept both boolean and string representations of true."""
    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):
    norm = np.linalg.norm(a) * np.linalg.norm(b)
    if norm == 0:
        return 0.0
    return float(np.dot(a, b) / norm)


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

def embed_query(query: str):
    emb = model.encode([query], prompt_name="query", normalize_embeddings=True)
    return round_floats(emb.tolist()[0])


def embed(text: str):
    emb = model.encode([text], normalize_embeddings=True)
    return round_floats(emb.tolist()[0])


def embedding(texts: list):
    embs = model.encode(texts, normalize_embeddings=True, batch_size=32)
    return round_floats(embs.tolist())


def similarity(text: str, texts: list):
    query_emb = model.encode([text], prompt_name="query", normalize_embeddings=True)[0]
    doc_embs = model.encode(texts, normalize_embeddings=True, batch_size=32)
    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")
    empty_cache()
    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})


@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")
    empty_cache()
    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})


@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")
    empty_cache()
    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})


@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")
    empty_cache()
    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': round_floats(result), 'took': took})


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


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