#!/usr/bin/env python3
"""A minimal JSON HTTP front for the question-answering pipeline, bound to loopback by default.

Purpose: put ask.py behind an endpoint so the service can be used from a script, from curl, or
    from the gateway in Part 9, without adding a web framework. It speaks JSON only: there is no
    HTML in this file, because the browser-facing front end for this project is Open WebUI,
    which already has accounts, TLS and a document feature. It refuses to start on a
    non-loopback address without an API key, because an unauthenticated endpoint on a network
    is an open text generator and a copy of everybody's questions.
Platform: all (standard library only, plus what ask.py needs)
Minimum memory: 8 GB on the machine running the models
Assumes: Python 3.9 or later, `sqlite-vec` and `pydantic` installed, ask.py in the same
    directory, an index built by ingest.py, and the servers ask.py needs.

Usage: python3 serve-qa.py --db qa-index.db --model qwen3-8b --port 8100
       LAN_IP=<this machine's address on your own network>; \
           python3 serve-qa.py --db qa-index.db --model qwen3-8b --host "$LAN_IP" \
           --port 8100 --api-key-file ~/.qa-key
"""

from __future__ import annotations

import argparse
import json
import sys
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

try:
    import ask
except ImportError:  # pragma: no cover - environment check
    sys.exit("ask.py must be in the same directory as this script; run it from there.")


USAGE = {
    "service": "private document question answering",
    "endpoints": {
        "GET /healthz": "readiness, and how many chunks are indexed",
        "POST /ask": 'send {"question": "..."} and receive the answer record',
    },
    "note": "JSON only. For a browser front end, use Open WebUI's Knowledge feature from Part 7.",
}

_local = threading.local()


def connection(db_path: str):
    """One SQLite connection per thread. SQLite connections are not shared across threads."""
    if getattr(_local, "db", None) is None:
        _local.db = ask.open_index(db_path)
    return _local.db


class Handler(BaseHTTPRequestHandler):
    server_version = "qa-service"
    args = None  # set in main()

    def _json(self, code: int, payload: dict) -> None:
        body = json.dumps(payload).encode("utf-8")
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _authorised(self) -> bool:
        # serve_key authenticates callers of THIS service. args.api_key is the separate
        # credential this service sends upstream to the model servers; do not confuse them.
        if not self.args.serve_key:
            return True
        header = self.headers.get("Authorization", "")
        return header == f"Bearer {self.args.serve_key}"

    def do_GET(self) -> None:  # noqa: N802 - name fixed by the base class
        if self.path.startswith("/healthz"):
            count = connection(self.args.db).execute("select count(*) from chunks").fetchone()[0]
            self._json(200, {"ok": True, "chunks": count, "model": self.args.model})
        elif self.path == "/":
            self._json(200, USAGE)
        else:
            self._json(404, {"error": "not found"})

    def do_POST(self) -> None:  # noqa: N802 - name fixed by the base class
        if not self.path.startswith("/ask"):
            self._json(404, {"error": "not found"})
            return
        if not self._authorised():
            self._json(401, {"error": "unauthorised"})
            return

        length = int(self.headers.get("Content-Length", "0") or 0)
        raw = self.rfile.read(length).decode("utf-8", "replace")
        try:
            question = json.loads(raw)["question"]
        except (ValueError, KeyError):
            self._json(400, {"error": 'send {"question": "..."}'})
            return
        self._json(200, ask.answer_question(question, connection(self.args.db), self.args))

    def log_message(self, fmt: str, *log_args) -> None:
        # Deliberately metadata only: the path and status, never the question. See the
        # security lesson in this part on what a model service's log contains by default.
        sys.stderr.write("%s - %s\n" % (self.address_string(), fmt % log_args))


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ask.add_common_arguments(parser)
    parser.add_argument("--host", default="127.0.0.1", help="bind address; loopback by default")
    parser.add_argument("--port", type=int, default=8100)
    parser.add_argument("--api-key-file", default=None,
                        help="file holding the bearer token callers must send")
    args = parser.parse_args()

    key = None
    if args.api_key_file:
        key = Path(args.api_key_file).expanduser().read_text(encoding="utf-8").strip() or None
    args.serve_key = key
    Handler.args = args

    loopback = args.host in {"127.0.0.1", "::1", "localhost"}
    if not loopback and not key:
        sys.exit("refusing to bind a non-loopback address without --api-key-file. "
                 "Read the security lesson in Part 10, then decide deliberately.")
    if not loopback:
        print(f"warning: binding {args.host}. Everything that can reach this address can ask "
              "your documents questions.", file=sys.stderr)

    connection(args.db)  # fail now if the index is missing, rather than on the first request
    server = ThreadingHTTPServer((args.host, args.port), Handler)
    print(f"listening on http://{args.host}:{args.port}/  (ctrl-c to stop)")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nstopped")
    finally:
        server.server_close()


if __name__ == "__main__":
    main()