from fastapi import FastAPI
from typing import Optional
import requests
import json
from pydantic import BaseModel


class Item(BaseModel):
    model: str
    messages: list
    temperature: float
    max_tokens: Optional[int] = None



# NOTE - we configure docs_url to serve the interactive Docs at the root path
# of the app. This way, we can use the docs as a landing page for the app on Spaces.
app = FastAPI()

API_ENDPOINT = "https://api.openai.com/v1/chat/completions"

def reog_file():
    with open('./openapis.txt', 'r+') as file:
        lines = file.readlines()
        lines.append(lines[0])
        lines = lines[1:]
        file.seek(0)
        file.writelines(lines)
        file.close()

def read_file():
    with open('./openapis.txt', 'r+') as file:
        lines = file.readlines()
        file.close()
    return lines, len(lines)

def generate_chat_completion(data):
    apis, num_apis = read_file()
    api_index = 0
    while(1):
        api_key = apis[api_index][:-1]
        print(api_key)
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}",
        }

        response = requests.post(API_ENDPOINT, headers=headers, data=data.json())
        if response.status_code == 200:
            return response.json()
        else:
            api_index += 1
            reog_file()
            if api_index == num_apis:
                raise Exception("Error")


@app.post("/generate")
def generate(data: Item):
    output = generate_chat_completion(data)
    return output