import http.server
import socketserver
import json
import sys

from chat_qwen import ChatQwen

bot = ChatQwen(sys.argv[1]) if len(sys.argv) > 1 else ChatQwen()


class Server(socketserver.TCPServer):
    # Avoid "address already used" error when frequently restarting the script 
    allow_reuse_address = True


class Handler(http.server.BaseHTTPRequestHandler):
    def do_GET(self):
        self.send_response(200)
        self.end_headers()
        self.wfile.write("Use POST with JSON body of the format \n"
                         "[{\"role\": \"user\", \"content\": \"message\"}] \n"
                         "or \n"
                         "[{\"role\": \"user\", \"content\": \"message\"}, "
                         "{\"role\": \"assistant\", \"content\": \"message\"}, "
                         "{\"role\": \"user\", \"content\": \"message\"}]".encode("utf-8"))

    def do_POST(self):
        try:
            content_len = int(self.headers.get('Content-Length'))
            post_body = self.rfile.read(content_len)
            json_body = json.loads(post_body)

            response = bot.generate(json_body)

            self.send_response(200)
            self.end_headers()
            self.wfile.write(json.dumps(response).encode("utf-8"))
        except Exception as e:
            print(e)
            self.send_response(400)
            self.end_headers()


with Server(("0.0.0.0", 8900), Handler) as httpd:
    httpd.serve_forever()