#!/usr/bin/env python3
import http.server
import urllib.request
import os

PORT = 8080
OLLAMA = "http://localhost:11434"
DIRECTORY = "/mnt/nvme/localai"

class Handler(http.server.SimpleHTTPRequestHandler):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, directory=DIRECTORY, **kwargs)

    def do_POST(self):
        if self.path.startswith("/ollama"):
            target = OLLAMA + self.path[7:]
            length = int(self.headers.get("Content-Length", 0))
            body = self.rfile.read(length)
            try:
                req = urllib.request.Request(
                    target, data=body,
                    headers={"Content-Type": "application/json"},
                    method="POST"
                )
                with urllib.request.urlopen(req) as resp:
                    self.send_response(200)
                    self.send_header("Content-Type", "application/x-ndjson")
                    self.send_header("Access-Control-Allow-Origin", "*")
                    self.end_headers()
                    while chunk := resp.read(512):
                        self.wfile.write(chunk)
                        self.wfile.flush()
            except Exception as e:
                self.send_error(502, str(e))

    def do_OPTIONS(self):
        self.send_response(200)
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "POST, GET, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type")
        self.end_headers()

    def log_message(self, format, *args):
        pass

if __name__ == "__main__":
    os.chdir(DIRECTORY)
    with http.server.ThreadingHTTPServer(("", PORT), Handler) as httpd:
        print(f"PaludarAI running on port {PORT}")
        httpd.serve_forever()