# Reranking using OpenAI Chat Completions API
#
# https://platform.openai.com/docs/guides/text-generation
#
# OpenAI has no native reranking endpoint. Relevance scores are obtained
# by prompting a chat model to score all documents against the query in
# a single call using JSON output mode.
#
# API compatible with rer1 (rer_qwen3_flask.py)
# POST: http://localhost:5007/rescore
# Request:  {"query": "search query", "docs": ["doc1 text", "doc2 text", ...]}
# Response: {"result": [0.92, 0.45, ...], "took": 123}
#
# Configuration via environment variables:
#   OPENAI_API_KEY    - required, your OpenAI API key
#   OPENAI_RER_MODEL  - optional, default: gpt-4o-mini
#   OPENAI_BASE_URL   - optional, default: https://api.openai.com/v1
#
import os
import json
import logging
import time

import openai
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_RER_MODEL", "gpt-4o-mini")
OPENAI_BASE_URL = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")

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

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

logger.info("OpenAI reranking service started")
logger.info("Model    : " + OPENAI_MODEL)
logger.info("Base URL : " + OPENAI_BASE_URL)

SYSTEM_PROMPT = (
    "You are a relevance scoring assistant. "
    "Given a search query and a list of documents, assign each document a relevance score between 0.01 and 1.0. "
    "Use the full range: "
    "0.9-1.0 = directly and fully answers the query; "
    "0.6-0.9 = highly relevant, addresses the main topic; "
    "0.3-0.6 = partially relevant, related topic or useful context; "
    "0.1-0.3 = loosely related, shares some terms or concepts; "
    "0.01-0.1 = very low relevance, almost unrelated. "
    "Never assign 0.0 — every document has at least minimal score. "
    "Spread scores across the range to differentiate documents. "
    "Output a JSON object with a single key \"scores\" containing an array of floats, "
    "one per document in the same order as the input. "
    "Return only the JSON object, no explanation."
)

def build_user_prompt(query, docs):
    lines = [f"Query: {query}", "", "Documents:"]
    for i, doc in enumerate(docs):
        lines.append(f"[{i}]: {doc}")
    lines.append("")
    lines.append(f'Return a JSON object: {{"scores": [score_0, score_1, ..., score_{len(docs)-1}]}}')
    return "\n".join(lines)

def rescore(query, docs):
    """
    Score each doc against the query using OpenAI chat completions.
    All documents are scored in a single API call for efficiency.
    Returns a list of float scores in the same order as docs.
    """
    user_prompt = build_user_prompt(query, docs)

    response = openai_client.chat.completions.create(
        model=OPENAI_MODEL,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user",   "content": user_prompt}
        ],
        response_format={"type": "json_object"},
        temperature=0.01,  # near-deterministic; avoids 403 on strict zero-temp restriction
        max_tokens=512
    )

    content = response.choices[0].message.content.strip()
#    logger.info("Raw response: " + content)

    try:
        parsed = json.loads(content)
        scores = parsed["scores"]
    except (json.JSONDecodeError, KeyError) as e:
        logger.error("Failed to parse OpenAI response: %s | content: %s", e, content)
        abort(502, description="Invalid response from OpenAI: " + str(e))

    # Ensure we have one score per document; fill missing with minimum score
    if len(scores) < len(docs):
        scores.extend([0.01] * (len(docs) - len(scores)))
    scores = [float(s) for s in scores[:len(docs)]]

    if len(set(scores)) == 1:
        logger.warning("All scores identical (%.2f) — reranker could not differentiate documents" % scores[0])

    logger.info("scores: " + str(scores))
    return scores


#====================================
# POST: http://localhost:8088/rescore
#{
#    "query": "train from London to Rome today",
#    "docs": [
#        "Schedule of transportation from London to Rome on train Schedule",
#        "Schedule of transportation from Rome to London on train and bus Schedule"
#    ]
#}

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

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

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