#!/usr/bin/env python3
"""OpenAI-compatible embeddings server for EmbeddingGemma 2 on an LLM-8850 (AXCL) card.

    python3 eg2_server.py --host 0.0.0.0 --port 8010

POST /v1/embeddings
    {"input": "text" | ["text", ...] | [{"text": "...", "image": "<base64 or data URL>", "audio": "<base64 wav>"}],
     "dimensions": 768|512|256|128, "encoding_format": "float"|"base64",
     "input_type": "query"|"document"|"sts"|"classification"|"clustering"|"code"}   (or "prompt_name")
GET  /v1/models, /health

Object inputs may hold several images/audios as lists; placeholders <|image|>/<|audio|> in "text"
mark where they go (otherwise they are appended). Local file paths in "image"/"audio" are only
accepted with --allow-local-files. One request runs at a time (the card is shared).
"""

import argparse
import base64
import io
import json
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

import numpy as np

from eg2_host import Eg2, load_audio

MODEL_ID = "embeddinggemma-2-axcl"


def decode_media(value, kind, allow_files):
    from PIL import Image

    if isinstance(value, str) and value.startswith("data:"):
        value = value.split(",", 1)[1]
    if isinstance(value, str) and allow_files and Path(value).exists():
        return Image.open(value) if kind == "image" else load_audio(value)
    raw = base64.b64decode(value)
    if kind == "image":
        return Image.open(io.BytesIO(raw))
    import soundfile as sf

    wav, sr = sf.read(io.BytesIO(raw), dtype="float32", always_2d=True)
    wav = wav.mean(1)
    if sr != 16000:
        from math import gcd

        from scipy.signal import resample_poly

        g = gcd(sr, 16000)
        wav = resample_poly(wav, 16000 // g, sr // g).astype(np.float32)
    return wav


def make_handler(eg, lock, allow_files):
    class Handler(BaseHTTPRequestHandler):
        def _send(self, code, obj):
            body = json.dumps(obj, ensure_ascii=False).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 do_GET(self):
            if self.path.rstrip("/") in ("/v1/models", "/models"):
                self._send(200, {"object": "list", "data": [{"id": MODEL_ID, "object": "model", "owned_by": "local"}]})
            elif self.path.rstrip("/") == "/health":
                self._send(200, {"status": "ok"})
            else:
                self._send(404, {"error": {"message": "not found"}})

        def do_POST(self):
            if self.path.rstrip("/") not in ("/v1/embeddings", "/embeddings"):
                return self._send(404, {"error": {"message": "not found"}})
            try:
                req = json.loads(self.rfile.read(int(self.headers.get("Content-Length", 0))) or b"{}")
                inputs = req.get("input")
                if isinstance(inputs, (str, dict)):
                    inputs = [inputs]
                if not inputs:
                    raise ValueError("input is required")
                dim = int(req.get("dimensions") or 768)
                prompt = req.get("input_type") or req.get("prompt_name")
                out, tokens = [], 0
                with lock:
                    for i, item in enumerate(inputs):
                        if isinstance(item, str):
                            v = eg.encode(text=item, prompt=prompt, dim=dim)
                        else:
                            imgs = item.get("image") or []
                            auds = item.get("audio") or []
                            imgs = [imgs] if isinstance(imgs, str) else imgs
                            auds = [auds] if isinstance(auds, str) else auds
                            v = eg.encode(text=item.get("text", ""),
                                          images=[decode_media(x, "image", allow_files) for x in imgs],
                                          audios=[decode_media(x, "audio", allow_files) for x in auds],
                                          prompt=prompt if item.get("text") else None, dim=dim)
                        tokens += eg.last_tokens
                        emb = (base64.b64encode(v.astype("<f4").tobytes()).decode()
                               if req.get("encoding_format") == "base64" else v.tolist())
                        out.append({"object": "embedding", "index": i, "embedding": emb})
                self._send(200, {"object": "list", "data": out, "model": MODEL_ID,
                                 "usage": {"prompt_tokens": tokens, "total_tokens": tokens}})
            except Exception as e:  # report bad requests as 400 with the reason
                self._send(400, {"error": {"message": f"{type(e).__name__}: {e}", "type": "invalid_request_error"}})

        def log_message(self, fmt, *args):
            print(time.strftime("%H:%M:%S"), self.address_string(), fmt % args)

    return Handler


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=8010)
    ap.add_argument("--assets", default=str(Path(__file__).resolve().parent / "assets"))
    ap.add_argument("--text-only", action="store_true", help="do not load the vision/audio encoders")
    ap.add_argument("--allow-local-files", action="store_true")
    args = ap.parse_args()
    eg = Eg2(args.assets, load=("text",) if args.text_only else ("text", "vision", "audio"))
    server = ThreadingHTTPServer((args.host, args.port), make_handler(eg, threading.Lock(), args.allow_local_files))
    print(f"serving {MODEL_ID} on http://{args.host}:{args.port}/v1/embeddings")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        eg.close()


if __name__ == "__main__":
    main()