import os
from http import HTTPStatus
from flask import Flask, request
from flask_caching import Cache
from flask_cors import CORS
import codegen
import gradio as gr



codgen_model = os.environ.get('MODEL', 'Salesforce/codegen-350M-mono')
print(f'Attempting to setup codegen with the model {codgen_model}')
codegen.setup(
    os.environ.get('MODEL', 'Salesforce/codegen-350M-mono'),
    setup_torch=os.environ.get("SETUP_TORCH") == 'true',
)


app = Flask(__name__)
CORS(app)
cache = Cache(app, config={'CACHE_TYPE': 'simple'})


# @app.route('/generate', methods=['GET'])
# @cache.cached(timeout=1440)  # Cache for 24 hours
def generate(user_input: str):
#    user_input = nl
    if not user_input:
        return 'Missing input', HTTPStatus.BAD_REQUEST
#    return {
#       'result': codegen.generate(user_input)
#   }, HTTPStatus.OK
    return codegen.generate(user_input)

# @app.route('/status', methods=['GET'])
def status():
    return {'status': 'OK'}
    
iface = gr.Interface(fn=generate, inputs="text", outputs="text")
iface.launch()