# LLM service using OpenAI Responses API (o4-mini)
#
# https://platform.openai.com/docs/guides/text-generation
#
# Endpoints: GET /run?task=..., POST /run, GET /health
#
# Supported request parameters:
#   task                 - required, the user prompt / instruction
#   prompt               - optional, overrides the system prompt for this request
#   max_tokens           - max output tokens (also accepted as max_output_tokens
#                          or max_new_tokens), default: 500
#   reasoning_effort     - reasoning depth: "minimal" | "medium", default: "medium"
#   verbosity            - answer length: "low" | "medium" | "high", optional
#   previous_response_id - optional, reuse prior reasoning context (stateful)
#
# Not supported by gpt-5 Responses API:
#   temperature, top_p, n, num_beams, early_stopping, max_time, min_tokens
#
# Configuration via environment variables:
#   OPENAI_API_KEY           - required, your OpenAI API key
#   OPENAI_LLM_MODEL         - optional, default: o4-mini
#   OPENAI_BASE_URL          - optional, default: https://api.openai.com/v1
#   OPENAI_SYSTEM_PROMPT     - optional, overrides the default system prompt
#   OPENAI_REASONING_EFFORT  - optional, default: medium
#   OPENAI_VERBOSITY         - optional, default: not set

import json
import logging
import os
import time

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_LLM_MODEL", "o4-mini")
OPENAI_BASE_URL         = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")
OPENAI_REASONING_EFFORT = os.environ.get("OPENAI_REASONING_EFFORT", "medium")
OPENAI_VERBOSITY        = os.environ.get("OPENAI_VERBOSITY", "")

DEFAULT_SYSTEM_PROMPT = (
    "You are a precise search assistant. "
    "Follow the instruction exactly. "
    "Return only the requested output — no explanations, no extra text, no markdown."
)
OPENAI_SYSTEM_PROMPT = os.environ.get("OPENAI_SYSTEM_PROMPT", DEFAULT_SYSTEM_PROMPT)

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 LLM service started")
logger.info("Model            : %s", OPENAI_MODEL)
logger.info("Base URL         : %s", OPENAI_BASE_URL)
logger.info("Reasoning effort : %s", OPENAI_REASONING_EFFORT)

# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------

def run(data: dict) -> str:
    task = (data.get("task") or "").strip()
    if not task:
        abort(400, description="Missing or empty 'task' field")

    max_output_tokens = int(
        data.get("max_output_tokens")
        or data.get("max_tokens")
        or data.get("max_new_tokens")
        or 500
    )
    reasoning_effort = data.get("reasoning_effort", OPENAI_REASONING_EFFORT)
    verbosity        = data.get("verbosity", OPENAI_VERBOSITY) or None
    system_prompt    = data.get("prompt") or data.get("system_prompt") or OPENAI_SYSTEM_PROMPT

    params = {
        "model":             OPENAI_MODEL,
        "max_output_tokens": max_output_tokens,
        "input": [
            {"role": "system", "content": system_prompt},
            {"role": "user",   "content": task},
        ],
        "reasoning": {"effort": reasoning_effort},
    }
    if verbosity:
        params["reasoning"]["summary"] = verbosity
    if "previous_response_id" in data:
        params["previous_response_id"] = data["previous_response_id"]

    response = client.responses.create(**params)

    result = response.output_text.strip()
    logger.info("Generated (%d chars): %r", len(result), result[:300])
    return result


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

# http://localhost:8087/run?task=Who+are+you%3F

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


@app.route("/run", methods=['GET'])
def run_get():
    task = (request.args.get('task') or '').strip()
    if not task:
        abort(400, description="Missing 'task' query parameter")
    logger.info("Task: %s", task[:200])
    try:
        start = time.perf_counter()
        result = run({"task": task})
        took = int((time.perf_counter() - start) * 1000)
        logger.info("Run 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("/run", methods=['POST'])
def run_post():
    try:
        data = json.loads(request.data)
    except (json.JSONDecodeError, ValueError):
        abort(400, description="Invalid JSON body")
    logger.info("Task: %s", str(data.get("task", ""))[:200])
    try:
        start = time.perf_counter()
        result = run(data)
        took = int((time.perf_counter() - start) * 1000)
        logger.info("Run 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("/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", debug=False)
