import base64
import io
import os
import sys
from typing import Any

import numpy as np
import torch
import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel
from PIL import Image

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from embedding.embedder import CompactEmbedder

app = FastAPI()
model = None


@app.on_event("startup")
def load_model():
    global model
    model = CompactEmbedder(local_dir="./models/multi-modal-embed-small")


class InlineData(BaseModel):
    mime_type: str
    data: str


class Part(BaseModel):
    text: str | None = None
    inline_data: InlineData | None = None


class Content(BaseModel):
    parts: list[Part]


class EmbedContentRequest(BaseModel):
    model: str = "models/multi-modal-embed-small"
    content: Content
    output_dimensionality: int | None = None
    task_type: str | None = None
    title: str | None = None


class BatchEmbedContentRequest(BaseModel):
    requests: list[EmbedContentRequest]


def _encode_content(content: Content) -> list[float]:
    vecs: list[Any] = []
    for part in content.parts:
        if part.text:
            vecs.append(model.encode_text(part.text))
        if part.inline_data:
            raw = base64.b64decode(part.inline_data.data)
            mime = part.inline_data.mime_type
            if mime.startswith("image/"):
                img = Image.open(io.BytesIO(raw)).convert("RGB")
                vecs.append(model.encode_image(img))
            elif mime.startswith("audio/"):
                audio = np.frombuffer(raw, dtype=np.float32)
                vecs.append(model.encode_audio(audio))
    if not vecs:
        vecs.append(model.encode_text(""))
    vec = vecs[0] if len(vecs) == 1 else torch.mean(torch.stack(vecs), dim=0)
    return vec.squeeze(0).tolist()


class EmbeddingResponse(BaseModel):
    embedding: dict[str, list[float]]


class EmbedContentResponse(BaseModel):
    embedding: dict[str, list[float]]


class BatchEmbedContentResponse(BaseModel):
    embeddings: list[dict[str, list[float]]]


@app.post("/v1beta/models/{model_id}:embedContent")
async def embed_content(
    model_id: str,
    req: EmbedContentRequest,
):
    emb = _encode_content(req.content)
    if req.output_dimensionality:
        emb = emb[: req.output_dimensionality]
    return {"embedding": {"values": emb}}


@app.post("/v1beta/models/{model_id}:batchEmbedContents")
async def batch_embed_contents(
    model_id: str,
    req: BatchEmbedContentRequest,
):
    embeddings = []
    for sub_req in req.requests:
        emb = _encode_content(sub_req.content)
        if sub_req.output_dimensionality:
            emb = emb[: sub_req.output_dimensionality]
        embeddings.append({"values": emb})
    return {"embeddings": embeddings}


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)