#!/usr/bin/env python3
"""Embedding proxy server backed by sagellm.

This server keeps the OpenAI-compatible embedding API surface (`/v1/embeddings`)
for existing callers, but delegates all embedding inference to a running sagellm
service instead of local torch/transformers inference.
"""

from __future__ import annotations

import argparse
import json
import logging
import time
from typing import Any
from urllib import error as urllib_error
from urllib import request as urllib_request

import uvicorn
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel

logger = logging.getLogger(__name__)
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)


class EmbeddingRequest(BaseModel):
    """OpenAI-compatible embedding request format."""

    input: str | list[str]
    model: str = "default"
    encoding_format: str = "float"


class EmbeddingProxyServer:
    """Proxy embeddings requests to sagellm service."""

    def __init__(self, model_name: str, upstream_base_url: str, api_key: str = "") -> None:
        self.model_name = model_name
        self.api_key = api_key
        self.upstream_base_url = self._normalize_base_url(upstream_base_url)

    @staticmethod
    def _normalize_base_url(base_url: str) -> str:
        normalized = base_url.rstrip("/")
        if normalized.endswith("/v1"):
            return normalized[:-3]
        return normalized

    def _headers(self) -> dict[str, str]:
        headers = {"Content-Type": "application/json"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        return headers

    def embed_texts(self, texts: list[str], model: str | None = None) -> list[list[float]]:
        payload: dict[str, Any] = {
            "model": model or self.model_name,
            "input": texts,
        }

        req = urllib_request.Request(
            url=f"{self.upstream_base_url}/v1/embeddings",
            data=json.dumps(payload).encode("utf-8"),
            headers=self._headers(),
            method="POST",
        )

        try:
            with urllib_request.urlopen(req, timeout=60.0) as resp:  # noqa: S310
                data = json.loads(resp.read().decode("utf-8"))
        except urllib_error.HTTPError as exc:
            body = exc.read().decode("utf-8", errors="replace")
            raise RuntimeError(f"sagellm HTTP {exc.code}: {body}") from exc
        except urllib_error.URLError as exc:
            raise RuntimeError(
                f"Cannot connect to sagellm at {self.upstream_base_url}: {exc.reason}"
            ) from exc

        items = data.get("data", [])
        return [item.get("embedding", []) for item in items]

    def health(self) -> bool:
        req = urllib_request.Request(
            url=f"{self.upstream_base_url}/health",
            headers=self._headers(),
            method="GET",
        )
        try:
            with urllib_request.urlopen(req, timeout=10.0):  # noqa: S310
                return True
        except Exception:
            return False


embedding_server: EmbeddingProxyServer | None = None

app = FastAPI(
    title="Embedding Proxy Server",
    description="OpenAI-compatible Embedding API proxy to sagellm",
    version="2.0",
)


@app.get("/")
async def root() -> dict[str, str]:
    return {
        "status": "ok",
        "backend": "sagellm",
        "model": embedding_server.model_name if embedding_server else "not loaded",
    }


@app.get("/health")
async def health() -> dict[str, str]:
    status = "ok" if embedding_server and embedding_server.health() else "degraded"
    return {
        "status": status,
        "backend": "sagellm",
        "model": embedding_server.model_name if embedding_server else "not loaded",
    }


@app.get("/v1/models")
async def list_models() -> dict[str, Any]:
    if not embedding_server:
        raise HTTPException(status_code=500, detail="Proxy server not initialized")

    return {
        "object": "list",
        "data": [
            {
                "id": embedding_server.model_name,
                "object": "model",
                "created": int(time.time()),
                "owned_by": "sagellm",
            }
        ],
    }


@app.post("/v1/embeddings")
async def create_embeddings(request: EmbeddingRequest) -> JSONResponse:
    if not embedding_server:
        raise HTTPException(status_code=500, detail="Proxy server not initialized")

    texts = [request.input] if isinstance(request.input, str) else request.input

    try:
        embeddings = embedding_server.embed_texts(texts, model=request.model)
    except Exception as exc:
        logger.exception("Embedding proxy failed: %s", exc)
        raise HTTPException(status_code=500, detail=str(exc)) from exc

    response_data = {
        "object": "list",
        "data": [
            {"object": "embedding", "embedding": emb, "index": idx}
            for idx, emb in enumerate(embeddings)
        ],
        "model": request.model,
        "usage": {
            "prompt_tokens": 0,
            "total_tokens": 0,
        },
    }
    return JSONResponse(content=response_data)


def main() -> None:
    parser = argparse.ArgumentParser(description="Embedding proxy server backed by sagellm")
    parser.add_argument("--model", type=str, default="BAAI/bge-m3")
    parser.add_argument("--port", type=int, default=8080)
    parser.add_argument("--host", type=str, default="0.0.0.0")
    parser.add_argument("--upstream", type=str, default="http://localhost:8000")
    parser.add_argument("--api-key", type=str, default="")
    parser.add_argument("--reload", action="store_true")
    args = parser.parse_args()

    global embedding_server
    embedding_server = EmbeddingProxyServer(
        model_name=args.model,
        upstream_base_url=args.upstream,
        api_key=args.api_key,
    )

    uvicorn.run(
        "embedding_server:app",
        host=args.host,
        port=args.port,
        reload=args.reload,
        log_level="info",
    )


if __name__ == "__main__":
    main()
