#!/usr/bin/env python3
"""Local parser service for Lain.

Serves one frozen filename parser over HTTP so a native or containerised Lain
can identify release names through a model. Read-only: never writes to disk and
holds exactly one artifact.

    GET  /health                 -> {"status": "ok", ...}
    POST /parse {"filename": …}  -> {"record": {...} | null, "spans": [...], ...}

The record is derived by the same gated contract resolver used at evaluation
time (``fp.normalize.gated_record``), so a null record is the model abstaining
and the caller should fall back to its own providers.

Usage:
    serve.py --artifact <prefix>/artifact/v3mix-bio-s1 --snapshot <prefix>/anifilebert
    serve.py --check "Some.Show.S01E02.1080p-GROUP.mkv"    # one-shot, prints JSON
"""

import argparse
import hashlib
import json
import logging
import os
import pathlib
import socket
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

import torch
from torch import nn

log = logging.getLogger("lain-parser")
APP = pathlib.Path(__file__).resolve().parent
MAX_BODY = 64 * 1024


class Parser:
    """One artifact, loaded once, guarded by a lock for the HTTP threads."""

    def __init__(self, artifact, snapshot):
        self.artifact = pathlib.Path(artifact).resolve()
        self.snapshot = pathlib.Path(snapshot).resolve()
        # The fp modules read the snapshot path at import time.
        os.environ["ANIFILEBERT_SNAPSHOT"] = str(self.snapshot)
        if str(APP) not in sys.path:
            sys.path.insert(0, str(APP))
        from fp import normalize
        from fp.parser2 import collate
        from fp.parser3_model import Parser as Model, decode_batch

        self.normalize, self.collate, self.decode_batch = normalize, collate, decode_batch
        config = json.loads((self.artifact / "config.json").read_text())
        model = Model(config["representation"], config["head"],
                      char_classes=config.get("char_classes", False)).eval()
        model = torch.ao.quantization.quantize_dynamic(model, {nn.Linear}, dtype=torch.qint8)
        state = torch.load(self.artifact / "model.int8.pt", map_location="cpu",
                           weights_only=False)
        model.load_state_dict(state)
        model.eval()
        self.model, self.config = model, config
        self.decoder = "viterbi" if config["head"] == "bio" else "span"
        self.limit = model.encoder.config.max_position_embeddings - 2
        self.lock = threading.Lock()
        self.warmup()

    def digest(self):
        return hashlib.sha256((self.artifact / "model.int8.pt").read_bytes()).hexdigest()

    @torch.inference_mode()
    def warmup(self):
        self.parse("Example.Show.S01E01.1080p-GROUP")

    @torch.inference_mode()
    def parse(self, query):
        """Return (record, spans, issues); record is None on abstention."""
        if len(query) > self.limit:
            return None, [], [], "over-encoder-limit"
        row = {"query": query, "record": None, "spans": [], "language": "real",
               "layout": "real", "view": "real", "family": "real"}
        with self.lock:
            item = self.model.encode(row)
            batch, _ = self.collate([item], self.model.encoder.config.pad_token_id)
            logits, chars = self.model(batch)
            spans = self.decode_batch(self.model, logits, chars, [item], self.decoder)[0]
        record, issues = self.normalize.gated_record(query, spans)
        return record, spans, issues, None


class Handler(BaseHTTPRequestHandler):
    server_version = "lain-parser/1.0"

    def log_message(self, fmt, *args):
        log.debug("%s %s", self.address_string(), fmt % args)

    def _send(self, code, payload):
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        self.send_response(code)
        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 do_GET(self):
        if self.path != "/health":
            self._send(404, {"error": "not-found"})
            return
        parser = self.server.parser
        self._send(200, {"status": "ok", "artifact": parser.artifact.name,
                         "head": parser.config["head"],
                         "max_chars": parser.limit,
                         "uptime_s": round(time.time() - self.server.started, 1)})

    def do_POST(self):
        if self.path != "/parse":
            self._send(404, {"error": "not-found"})
            return
        try:
            length = int(self.headers.get("Content-Length") or 0)
        except ValueError:
            length = 0
        if length <= 0 or length > MAX_BODY:
            self._send(400, {"error": "empty-body"})
            return
        try:
            payload = json.loads(self.rfile.read(length))
        except (ValueError, UnicodeDecodeError):
            self._send(400, {"error": "invalid-json"})
            return
        filename = payload.get("filename") if isinstance(payload, dict) else None
        if not isinstance(filename, str) or not filename.strip():
            self._send(400, {"error": "invalid-filename"})
            return
        began = time.perf_counter()
        try:
            record, spans, issues, limited = self.server.parser.parse(filename)
        except ValueError as exc:
            self._send(400, {"error": "invalid-filename", "detail": str(exc)})
            return
        except Exception:  # never take the service down over one bad request
            log.exception("parse failed for %r", filename[:120])
            self._send(500, {"error": "internal"})
            return
        elapsed = round(1000 * (time.perf_counter() - began), 2)
        log.info("parse %.2fms abstain=%s chars=%d", elapsed, record is None, len(filename))
        self._send(200, {"artifact": self.server.parser.artifact.name,
                         "record": record, "spans": spans, "issues": issues,
                         "abstain_reason": limited, "latency_ms": elapsed})


def build_parsers(artifact, snapshot):
    parser = Parser(artifact, snapshot)
    log.info("artifact %s head=%s max_chars=%d sha256=%s", parser.artifact.name,
             parser.config["head"], parser.limit, parser.digest()[:16])
    return parser


class UnixHTTPServer(ThreadingHTTPServer):
    """HTTP over a Unix domain socket, for containers that cannot reach the host.

    A container with Lain's data directory bind-mounted sees the socket at
    /data/parser.sock and needs no firewall rule or published port.
    """

    address_family = socket.AF_UNIX
    daemon_threads = True

    def server_bind(self):
        path = self.server_address
        if os.path.exists(path):
            os.unlink(path)
        self.socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.socket.bind(path)
        os.chmod(path, 0o666)  # the container's user is not our user
        self.server_name = "lain-parser"

    def server_activate(self):
        self.socket.listen(self.request_queue_size)

    def get_request(self):
        request, _ = self.socket.accept()
        return request, ("unix", 0)

    def server_close(self):
        super().server_close()
        try:
            os.unlink(self.server_address)
        except OSError:
            pass


def serve(hosts, port, parser, unix_path):
    started = time.time()
    servers = []
    for host in hosts:
        try:
            server = ThreadingHTTPServer((host, port), Handler)
        except OSError as exc:
            log.warning("cannot bind %s:%d: %s", host, port, exc)
            continue
        server.parser = parser
        server.started = started
        servers.append(server)
    if unix_path:
        try:
            server = UnixHTTPServer(unix_path, Handler)
        except OSError as exc:
            log.warning("cannot bind unix socket %s: %s", unix_path, exc)
        else:
            server.parser = parser
            server.started = started
            servers.append(server)
    if not servers:
        raise SystemExit("no bind address worked; is the port already in use?")
    for server in servers:
        threading.Thread(target=server.serve_forever, daemon=True).start()
        if isinstance(server, UnixHTTPServer):
            log.info("listening on unix://%s", server.server_address)
        else:
            log.info("listening on http://%s:%d", *server.server_address[:2])
    try:
        while True:
            time.sleep(3600)
    except KeyboardInterrupt:
        log.info("shutting down")


def main(argv=None):
    args = parse_args(argv)
    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO,
                        format="%(asctime)s %(levelname)s %(message)s")
    parser = build_parsers(args.artifact, args.snapshot)
    if args.check is not None:
        record, spans, issues, limited = parser.parse(args.check)
        print(json.dumps({"filename": args.check, "record": record, "spans": spans,
                          "issues": issues, "abstain_reason": limited},
                         ensure_ascii=False, indent=2))
        return 0
    serve([h.strip() for h in args.host.split(",") if h.strip()], args.port, parser,
          args.unix_socket)
    return 0


def parse_args(argv):
    default_prefix = pathlib.Path(os.environ.get("LAIN_PARSER_PREFIX",
                                                 pathlib.Path.home() / ".local/share/lain-parser"))
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--artifact", default=str(default_prefix / "artifact/v3mix-bio-s1"))
    ap.add_argument("--snapshot", default=str(default_prefix / "anifilebert"))
    ap.add_argument("--host", default=os.environ.get("LAIN_PARSER_HOST", "127.0.0.1,172.17.0.1"),
                    help="comma-separated bind addresses")
    ap.add_argument("--port", type=int, default=int(os.environ.get("LAIN_PARSER_PORT", "9377")))
    ap.add_argument("--unix-socket", default=os.environ.get("LAIN_PARSER_SOCKET", ""),
                    help="also listen on this Unix socket path")
    ap.add_argument("--check", metavar="FILENAME",
                    help="parse one filename, print JSON and exit")
    ap.add_argument("--verbose", action="store_true")
    return ap.parse_args(argv)


if __name__ == "__main__":
    sys.exit(main())