from http.server import HTTPServer, SimpleHTTPRequestHandler
import urllib.request, json, os

class Handler(SimpleHTTPRequestHandler):
    def __init__(self, *a, **kw):
        super().__init__(*a, directory='/mnt/ssd', **kw)
    def do_POST(self):
        if self.path == '/proxy/chat':
            length = int(self.headers['Content-Length'])
            body = self.rfile.read(length)
            req = urllib.request.Request('http://localhost:8080/v1/chat/completions', body, {'Content-Type':'application/json'})
            try:
                with urllib.request.urlopen(req, timeout=300) as r:
                    data = r.read()
                self.send_response(200)
                self.send_header('Content-Type','application/json')
                self.end_headers()
                self.wfile.write(data)
            except Exception as e:
                self.send_response(502)
                self.end_headers()
                self.wfile.write(json.dumps({'error':str(e)}).encode())
        else:
            self.send_response(404)
            self.end_headers()
    def do_GET(self):
        if self.path == '/proxy/models':
            try:
                with urllib.request.urlopen('http://localhost:8080/v1/models', timeout=10) as r:
                    data = r.read()
                self.send_response(200)
                self.send_header('Content-Type','application/json')
                self.end_headers()
                self.wfile.write(data)
            except:
                self.send_response(502)
                self.end_headers()
        else:
            super().do_GET()

print('Serving on :9090')
HTTPServer(('0.0.0.0',9090),Handler).serve_forever()