import requests
import json
from io import StringIO
import sys
import gradio
import os


def generate_response(query, top_p):
    endpoint = os.environ.get("ENDPOINT", None)
    token = os.environ.get("TOKEN", None)
    params = {
        'accept': 'application/json',
        'Content-Type': 'application/json',
        "Authorization": f"Bearer {token}"
    }
    data = {
        'inputs': query,
        "parameters": {
            "details": True,
            "do_sample": True,
            "max_new_tokens": 256,
            "repetition_penalty": 1.2,
            "seed": 0,
            "stop": ['<|endoftext|>'],
            "temperature": 0.2,
            "top_p": top_p,
        }
    }
    try:
        res = requests.post(endpoint, data=json.dumps(data), headers=params)
        res = json.loads(res.content.decode("utf-8"))
        return res[0]["generated_text"].split("<|endoftext|>")[0].strip()
    except:
        raise gradio.Error("Connection Error. Please check your ACCESS_TOKEN")


def run_code(code):
    output_buffer = StringIO()
    sys.stdout = output_buffer
    exec(code, {})
    sys.stdout = sys.__stdout__
    pred = output_buffer.getvalue().strip()
    return pred