# OpenAI-compatible Image Generation API
# POST /v1/image/generations
# Required params: model, catelogid
# Optional params: prompt, image_url, n, size, style, quality, etc.
#
# Example request:
# {
#     "model": "jimeng-4.0",
#     "catelogid": "t2i",
#     "prompt": "A beautiful sunset over the ocean",
#     "size": "1024x1024",
#     "n": 1
# }
#
# Response format depends on the upstream model (sync returns image data, async returns task info)

import json
import time
from functools import partial
from appPublic.log import debug
from appPublic.dictObject import DictObject
from appPublic.uniqueID import getID
from appPublic.timeUtils import curDateString, timestampstr
from sqlor.dbpools import get_sor_context

debug_params('params_kw', params_kw)

userid = await get_user()
userorgid = await get_userorgid()
if userid is None:
    debug('need login')
    return openai_403()

# Validate required parameters
if not params_kw.model:
    d = return_error('Missing required parameter: model')
    return json_response(d, status=400)

if not params_kw.catelogid:
    d = return_error('Missing required parameter: catelogid')
    return json_response(d, status=400)

if not params_kw.prompt:
    d = return_error('Missing required parameter: prompt')
    return json_response(d, status=400)

lctype = params_kw.catelogid

env = request._run_ns
async with get_sor_context(env, 'llmage') as sor:
    # Look up llm by model name and catalog type through llm_api_map
    sql = """select distinct a.* from llm a
join llm_api_map m on a.id = m.llmid
join llmcatelog b on m.llmcatelogid = b.id
where (b.id = ${lctype}$ OR b.name = ${lctype}$)
    and a.model=${model}$
    and a.status = 'published'"""
    recs = await sor.sqlExe(sql, {
        'lctype': lctype,
        'model': params_kw.model
    })
    if len(recs) == 0:
        debug(f'{params_kw.model=} not found for catalog {lctype}')
        return openai_400()
    params_kw.llmid = recs[0].id

debug(f'{params_kw.llmid=}')

# Check balance
f = await checkCustomerBalance(params_kw.llmid, userid, userorgid)
if not f:
    debug(f'{userid=} balance not enough')
    return openai_429()

# Generate task ID and attach to params
if not params_kw.transno:
    params_kw.transno = getID()

# Call inference (image generation can be sync or async depending on model config)
return await inference(request, env=env)