from __future__ import annotations

import json
import socket
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

import ollama


ROOT = Path(__file__).parent
MODEL = "qwen2.5:3b"
HOST = "0.0.0.0"
PORT = 8787


class Handler(SimpleHTTPRequestHandler):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, directory=str(ROOT), **kwargs)

    def do_POST(self) -> None:
        if self.path != "/api/chat":
            self.send_error(404)
            return

        length = int(self.headers.get("content-length", 0))
        data = json.loads(self.rfile.read(length) or b"{}")
        messages = data.get("messages", [])

        try:
            response = ollama.chat(
                model=MODEL,
                messages=[
                    {
                        "role": "system",
                        "content": "You are a concise local assistant. Reply in Korean unless the user asks otherwise.",
                    },
                    *messages,
                ],
            )
            payload = {"reply": response["message"]["content"]}
            self._send_json(200, payload)
        except Exception as exc:
            self._send_json(500, {"error": str(exc)})

    def _send_json(self, status: int, payload: dict[str, str]) -> None:
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("content-type", "application/json; charset=utf-8")
        self.send_header("content-length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)


def main() -> None:
    server = ThreadingHTTPServer((HOST, PORT), Handler)
    local_ip = socket.gethostbyname(socket.gethostname())
    print(f"LLM web is running locally at http://127.0.0.1:{PORT}")
    print(f"LAN access is available at http://{local_ip}:{PORT}")
    server.serve_forever()


if __name__ == "__main__":
    main()