import base64
import binascii
import json
import mimetypes
import secrets
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from io import BytesIO
from pathlib import Path
from urllib.parse import urlparse

from generate import CHECKPOINTS, MODEL_ID, MODEL_SIZES, generate_images, is_model_loaded, load_model


ROOT = Path(__file__).resolve().parent


def encode_frame(image):
    buffer = BytesIO()
    image.save(buffer, format="JPEG", quality=86)
    return base64.b64encode(buffer.getvalue()).decode("ascii")


def generate_frames(seed, checkpoints=CHECKPOINTS, initial_image=None, model_id=MODEL_ID):
    yield {"type": "status", "status": "loading", "model": model_id}

    try:
        for step, image in generate_images(seed, checkpoints=checkpoints, initial_image=initial_image, model_id=model_id):
            yield {"type": "progress", "step": step, "total": 1000}
            if image is not None:
                yield {"type": "frame", "step": step, "image": encode_frame(image)}

        _, _, device = load_model(model_id=model_id)
        yield {"type": "complete", "seed": seed, "device": device, "model": model_id}
    except Exception as error:
        yield {"type": "error", "message": str(error)}


class DiffusionHandler(BaseHTTPRequestHandler):
    def send_bytes(self, status, content_type, body):
        self.send_response(status)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):
        path = urlparse(self.path).path
        if path == "/api/health":
            body = json.dumps({"model": MODEL_ID, "models": list(MODEL_SIZES), "loaded": is_model_loaded()}).encode()
            self.send_bytes(200, "application/json", body)
            return

        relative_path = "index.html" if path == "/" else path.lstrip("/")
        if relative_path not in {
            "index.html",
            "app.js",
            "styles.css",
            "assets/celeba_hq_validation_91_256.png",
            "assets/cifar10_train_23_32.png",
        }:
            self.send_error(404)
            return

        file_path = ROOT / relative_path
        content_type = mimetypes.guess_type(file_path.name)[0] or "application/octet-stream"
        self.send_bytes(200, content_type, file_path.read_bytes())

    def do_POST(self):
        if urlparse(self.path).path != "/api/generate":
            self.send_error(404)
            return

        try:
            length = int(self.headers.get("Content-Length", "0"))
            request = json.loads(self.rfile.read(length) or b"{}")
            model_id = request.get("model", MODEL_ID)
            if not isinstance(model_id, str) or model_id not in MODEL_SIZES:
                raise ValueError("Unsupported model.")
            seed = int(request.get("seed", secrets.randbelow(2**31)))
            if seed < 0 or seed >= 2**63:
                raise ValueError("Seed must be between 0 and 2^63 - 1.")
            checkpoints = request.get("checkpoints", CHECKPOINTS)
            if (
                not isinstance(checkpoints, list)
                or not 2 <= len(checkpoints) <= 12
                or any(type(step) is not int or step < 0 or step > 1000 for step in checkpoints)
            ):
                raise ValueError("Checkpoints must be a list of 2 to 12 integers from 0 to 1000.")
            checkpoints = tuple(sorted(set(checkpoints)))
            if checkpoints[0] != 0 or checkpoints[-1] != 1000:
                raise ValueError("Checkpoints must include 0 and 1000.")
            initial_image = None
            encoded_image = request.get("initial_image")
            if encoded_image is not None:
                if not isinstance(encoded_image, str) or len(encoded_image) > 3_000_000:
                    raise ValueError("The forward-noise image is invalid or too large.")
                from PIL import Image

                encoded_image = encoded_image.partition(",")[2]
                initial_image = Image.open(BytesIO(base64.b64decode(encoded_image, validate=True)))
                initial_image.load()
        except (ValueError, TypeError, json.JSONDecodeError, binascii.Error, OSError) as error:
            self.send_bytes(400, "application/json", json.dumps({"error": str(error)}).encode())
            return

        self.send_response(200)
        self.send_header("Content-Type", "application/x-ndjson; charset=utf-8")
        self.send_header("Cache-Control", "no-store")
        self.send_header("X-Accel-Buffering", "no")
        self.end_headers()
        try:
            for event in generate_frames(seed, checkpoints, initial_image, model_id):
                self.wfile.write((json.dumps(event) + "\n").encode())
                self.wfile.flush()
        except (BrokenPipeError, ConnectionResetError):
            pass

    def log_message(self, format_string, *args):
        print(f"{self.address_string()} - {format_string % args}")


class DiffusionServer(ThreadingHTTPServer):
    daemon_threads = True


if __name__ == "__main__":
    server = DiffusionServer(("0.0.0.0", 8000), DiffusionHandler)
    print(f"Diffusion Lab: http://{server.server_address[0]}:{server.server_address[1]}")
    print(f"Model: {MODEL_ID} (loaded on first generation request)")
    server.serve_forever()