#!/usr/bin/env python3
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.

# Long-lived daemon that keeps ONE SDXL model resident in XPU VRAM so prompt
# iteration skips the model-reload tax. Stdlib HTTP only. Serial by design:
# one GPU, so concurrent generation would OOM. See README.md.

import io
import json
import os
import sys
from http.server import BaseHTTPRequestHandler, HTTPServer

import torch

import generate

PORT = int(os.environ.get("IMGGEN_PORT", "8765"))

STATE = {"model": None, "pipe": None}


def load(model_key):
    if STATE["model"] == model_key and STATE["pipe"] is not None:
        return
    if STATE["pipe"] is not None:
        del STATE["pipe"]
        STATE["pipe"] = None
        STATE["model"] = None
        torch.xpu.empty_cache()
    STATE["pipe"] = generate.build_pipe(model_key)
    STATE["model"] = model_key


class Handler(BaseHTTPRequestHandler):
    def _json(self, code, obj):
        body = json.dumps(obj).encode()
        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 _read_json(self):
        length = int(self.headers.get("Content-Length", "0"))
        raw = self.rfile.read(length) if length else b"{}"
        return json.loads(raw)  # raises on bad JSON -> caught by caller

    def do_GET(self):
        if self.path == "/status":
            self._json(200, {"model": STATE["model"], "ready": STATE["pipe"] is not None})
        else:
            self._json(404, {"error": "not found"})

    def do_POST(self):
        try:
            data = self._read_json()
        except (ValueError, json.JSONDecodeError):
            self._json(400, {"error": "bad JSON"})
            return

        if self.path == "/model":
            name = data.get("name")
            try:
                generate.resolve_model(name)
            except SystemExit as e:
                self._json(400, {"error": str(e)})
                return
            try:
                load(name)
            except Exception as e:  # download / load failure
                self._json(500, {"error": f"load failed: {e}"})
                return
            self._json(200, {"model": STATE["model"], "ready": True})
            return

        if self.path == "/generate":
            try:
                image, elapsed = generate.generate(
                    STATE["pipe"],
                    data["prompt"],
                    data.get("negative", ""),
                    int(data.get("steps", 25)),
                    float(data.get("guidance", 7.0)),
                    data.get("seed"),
                    int(data.get("width", 1024)),
                    int(data.get("height", 1024)),
                )
            except torch.OutOfMemoryError:
                torch.xpu.empty_cache()
                self._json(500, {"error": "OOM, try fewer steps or smaller size"})
                return
            except Exception as e:
                self._json(500, {"error": str(e)})
                return
            buf = io.BytesIO()
            image.save(buf, format="PNG")
            body = buf.getvalue()
            self.send_response(200)
            self.send_header("Content-Type", "image/png")
            self.send_header("X-Elapsed", f"{elapsed:.2f}")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
            return

        self._json(404, {"error": "not found"})

    def log_message(self, fmt, *args):
        sys.stderr.write("%s - %s\n" % (self.address_string(), fmt % args))


def main():
    startup = os.environ.get("IMGGEN_MODEL", generate.DEFAULT_MODEL)
    print(f"Loading {startup} on XPU ...", file=sys.stderr)
    load(startup)
    print(f"Ready on 127.0.0.1:{PORT} with model {startup}", file=sys.stderr)
    HTTPServer(("0.0.0.0", PORT), Handler).serve_forever()


if __name__ == "__main__":
    main()