# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.

# SPDX-License-Identifier: Apache-2.0

import asyncio
import base64
import logging
import os
import sys
import time
from contextlib import asynccontextmanager
from typing import Any, Optional

import numpy as np
import orjson
from fastapi import FastAPI, Request
from fastapi.concurrency import run_in_threadpool
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse, Response
from pydantic import BaseModel, StrictInt
from starlette.exceptions import HTTPException as StarletteHTTPException

from models.autoports.google_embeddinggemma_2.server.engine import (
    DIMENSIONS,
    MODEL_ALIASES,
    SERVED_MODEL,
    RequestError,
    build_engine,
    normalize_inputs,
    shutdown_engine,
)

log = logging.getLogger("embeddinggemma2.server")
INLINE_TEXTS = 8


class EmbeddingRequest(BaseModel):
    input: Any = None
    model: Optional[str] = None
    dimensions: Optional[StrictInt] = None
    encoding_format: Optional[str] = "float"
    input_type: Optional[str] = None
    user: Optional[str] = None


def error_response(status, message, param=None, code=None, kind="invalid_request_error"):
    return JSONResponse(
        status_code=status, content={"error": {"message": message, "type": kind, "param": param, "code": code}}
    )


@asynccontextmanager
async def lifespan(app):
    t0 = time.perf_counter()
    engine = build_engine()
    built = time.perf_counter()
    warmed = engine.warm()
    engine.startup_s = {"build_and_capture": round(built - t0, 1), "warm": round(time.perf_counter() - built, 1)}
    if DEMO:
        from models.autoports.google_embeddinggemma_2.demo.mount import build_demo_index

        t1 = time.perf_counter()
        await run_in_threadpool(build_demo_index, app, engine)
        engine.startup_s["demo_index"] = round(time.perf_counter() - t1, 1)
    app.state.engine = engine
    log.warning(
        "embeddinggemma-2 ready: policy=%s mesh=%s grid=%s slots=%s warmed=%d startup=%s",
        engine.policy_id,
        engine.mesh_shape,
        engine.grid,
        engine.slot_summary(),
        warmed,
        engine.startup_s,
    )
    try:
        yield
    finally:
        app.state.engine = None
        shutdown_engine(engine)


DEMO = os.environ.get("EMBEDDINGGEMMA_DEMO", "0") == "1"
app = FastAPI(title="EmbeddingGemma 2 on Tenstorrent p150", lifespan=lifespan)


@app.exception_handler(RequestError)
async def request_error_handler(request: Request, exc: RequestError):
    return error_response(exc.status, exc.message, exc.param, exc.code)


@app.exception_handler(RequestValidationError)
async def validation_error_handler(request: Request, exc: RequestValidationError):
    parts = []
    for e in exc.errors():
        loc = ".".join(str(x) for x in e.get("loc", ()) if x != "body")
        parts.append(f"{loc}: {e.get('msg')}" if loc else str(e.get("msg")))
    return error_response(400, "; ".join(parts) or "invalid request")


@app.exception_handler(StarletteHTTPException)
async def http_error_handler(request: Request, exc: StarletteHTTPException):
    return error_response(
        exc.status_code, str(exc.detail), kind="not_found_error" if exc.status_code == 404 else "error"
    )


@app.exception_handler(Exception)
async def internal_error_handler(request: Request, exc: Exception):
    log.exception("unhandled error")
    return error_response(500, f"internal error: {type(exc).__name__}: {exc}", kind="server_error")


def engine_or_503(request):
    engine = getattr(request.app.state, "engine", None)
    if engine is None:
        raise RequestError("model is not loaded", status=503, code="not_ready")
    if not engine.alive():
        raise RequestError("device thread is not running", status=503, code="device_thread_dead")
    return engine


def health_payload(engine):
    return {
        "status": "ok",
        "model": SERVED_MODEL,
        "precision_policy": engine.policy_id,
        "mesh_shape": engine.mesh_shape,
        "grid": engine.grid,
        "max_tokens": engine.max_tokens,
        "switch_interval_ms": round(sys.getswitchinterval() * 1000, 3),
        "batching": {
            "max_batch_items": engine.max_batch_items,
            "batch_wait_ms": engine.batch_wait_s * 1000,
            "slot_ms": {f"{b}x{p}": round(t * 1e3, 2) for (b, p), t in sorted(engine.slot_cost.items())},
        },
        "dimensions": list(DIMENSIONS),
        "slots": engine.slot_summary(),
        "startup_s": engine.startup_s,
        "queue_depth": engine.queue_depth(),
        "stats": dict(engine.stats),
        "device_thread_alive": engine.alive(),
        "bucket_grouping": engine.bucket_grouping,
        "traces": engine.trace_counters(),
        "uptime_s": round(time.time() - engine.started, 1),
    }


@app.get("/health")
async def health(request: Request):
    return health_payload(engine_or_503(request))


@app.get("/v1/models")
async def models(request: Request):
    engine = engine_or_503(request)
    return {
        "object": "list",
        "data": [
            {
                "id": SERVED_MODEL,
                "object": "model",
                "created": int(engine.started),
                "owned_by": "google",
                "root": SERVED_MODEL,
                "max_model_len": engine.max_tokens,
                "dimensions": list(DIMENSIONS),
                "modalities": ["text", "image", "audio", "video"],
            }
        ],
    }


def encode_vector(vec, fmt):
    vec = np.ascontiguousarray(vec, dtype="<f4")
    if fmt == "base64":
        return base64.b64encode(vec.tobytes()).decode("ascii")
    return vec


@app.post("/v1/embeddings")
async def embeddings(body: EmbeddingRequest, request: Request):
    engine = engine_or_503(request)
    if body.model is not None and body.model not in MODEL_ALIASES:
        raise RequestError(f"model {body.model!r} is not served; use {SERVED_MODEL!r}", "model", 404, "model_not_found")
    dim = 768 if body.dimensions is None else body.dimensions
    if dim not in DIMENSIONS:
        raise RequestError(f"dimensions must be one of {list(DIMENSIONS)}, got {dim}", "dimensions")
    fmt = body.encoding_format or "float"
    if fmt not in ("float", "base64"):
        raise RequestError(f"encoding_format must be 'float' or 'base64', got {fmt!r}", "encoding_format")
    if body.input is None:
        raise RequestError("input is required", "input")
    engine.prompt_for(body.input_type)
    raw = body.input
    if isinstance(raw, str) or (
        isinstance(raw, list) and len(raw) <= INLINE_TEXTS and all(isinstance(x, str) for x in raw)
    ):
        items = engine.prepare(normalize_inputs(raw), body.input_type)
    else:
        inputs = await run_in_threadpool(normalize_inputs, raw)
        items = await run_in_threadpool(engine.prepare, inputs, body.input_type)
    vectors = await asyncio.wrap_future(engine.submit(items, dim))
    engine.stats["requests"] += 1
    engine.stats["inputs"] += len(items)
    tokens = sum(len(it["input_ids"]) for it in items)
    payload = {
        "object": "list",
        "data": [
            {"object": "embedding", "index": i, "embedding": encode_vector(v, fmt)} for i, v in enumerate(vectors)
        ],
        "model": SERVED_MODEL,
        "usage": {"prompt_tokens": tokens, "total_tokens": tokens},
    }
    return Response(orjson.dumps(payload, option=orjson.OPT_SERIALIZE_NUMPY), media_type="application/json")


if DEMO:
    from models.autoports.google_embeddinggemma_2.demo.mount import mount_demo

    mount_demo(app, lambda a: health_payload(a.state.engine) if getattr(a.state, "engine", None) else {})