import os
import json
import requests
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import parse_qs, urlparse

from inference import t5_infer

# https://huggingface.co/settings/tokens
# https://huggingface.co/spaces/{username}/{space}/settings
API_TOKEN = os.getenv('BIG_GAN_TOKEN')

class RequestHandler(SimpleHTTPRequestHandler):
    def do_GET(self):
        if self.path == "/":
            self.path = "index.html"

            return SimpleHTTPRequestHandler.do_GET(self)

        if self.path.startswith("/biggan_infer"):
            input = parse_qs(urlparse(self.path).query).get("input", None)[0]

            output = requests.request(
                "POST",
                "https://api-inference.huggingface.co/models/osanseviero/BigGAN-deep-128",
                headers={"Authorization": f"Bearer {API_TOKEN}"},
                data=json.dumps(input)
            )

            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.end_headers()

            self.wfile.write(output.content)

            return SimpleHTTPRequestHandler

        elif self.path.startswith("/t5_infer"):
            input = parse_qs(urlparse(self.path).query).get("input", None)

            output = t5_infer(input)

            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.end_headers()

            self.wfile.write(json.dumps({"output": output}).encode("utf-8"))

            return SimpleHTTPRequestHandler

        else:
            return SimpleHTTPRequestHandler.do_GET(self)


server = ThreadingHTTPServer(("", 7860), RequestHandler)

server.serve_forever()