"""Production-oriented text-and-still-image API for Phillnet Mini Text-Vision.

The service intentionally exposes only chat completion and visual-question-answering
workflows. SDXL, image/video synthesis, audio, agents, tools, and remote image URL
fetching are outside this deployment surface.
"""
from __future__ import annotations

import base64
import hmac
import io
import os
import threading
import time
import uuid
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any

import torch
from fastapi import Depends, FastAPI, Header, HTTPException, Request, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from PIL import Image, UnidentifiedImageError
from pydantic import BaseModel, Field
from transformers import AutoModelForCausalLM, AutoProcessor

MODEL_DIR = Path(os.getenv("MODEL_DIR", Path(__file__).resolve().parent))
API_KEY = os.getenv("PHILLNET_API_KEY", "")
MAX_IMAGE_BYTES = int(os.getenv("PHILLNET_MAX_IMAGE_BYTES", str(10 * 1024 * 1024)))
MAX_IMAGE_PIXELS = int(os.getenv("PHILLNET_MAX_IMAGE_PIXELS", str(24_000_000)))
MAX_REQUEST_BYTES = int(os.getenv("PHILLNET_MAX_REQUEST_BYTES", str(12 * 1024 * 1024)))
CORS_ORIGINS = [origin.strip() for origin in os.getenv("PHILLNET_CORS_ORIGINS", "").split(",") if origin.strip()]

Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS
MODEL: Any | None = None
PROCESSOR: Any | None = None
GENERATION_LOCK = threading.Lock()


class ImageContent(BaseModel):
    type: str
    text: str | None = Field(default=None, max_length=32_000)
    image_base64: str | None = None


class Message(BaseModel):
    role: str = Field(pattern="^(system|user|assistant)$")
    content: str | list[ImageContent]


class ChatRequest(BaseModel):
    model: str = "phillnet-mini-text-vision"
    messages: list[Message] = Field(min_length=1, max_length=32)
    max_tokens: int = Field(default=8192, ge=1, le=8192)
    temperature: float = Field(default=0.0, ge=0.0, le=2.0)
    reasoning_effort: str = Field(default="max", pattern="^(direct|low|medium|high|max)$")


def require_api_key(
    authorization: str | None = Header(default=None),
    x_api_key: str | None = Header(default=None),
) -> None:
    """Enforce an API key when PHILLNET_API_KEY is configured.

    Local development remains frictionless when the environment variable is empty.
    Production compose configuration supplies a non-empty secret by default.
    """
    if not API_KEY:
        return
    candidate = x_api_key or ""
    if authorization and authorization.lower().startswith("bearer "):
        candidate = authorization[7:].strip()
    if not candidate or not hmac.compare_digest(candidate, API_KEY):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Valid API credentials are required.",
            headers={"WWW-Authenticate": "Bearer"},
        )


def decode_image(encoded: str) -> Image.Image:
    try:
        raw = encoded.split(",", 1)[1] if encoded.startswith("data:") else encoded
        # Base64 expands bytes by roughly 4/3. Guard before decoding a large body.
        if len(raw) > ((MAX_IMAGE_BYTES * 4) // 3) + 8:
            raise HTTPException(413, f"image_base64 exceeds the {MAX_IMAGE_BYTES}-byte limit")
        payload = base64.b64decode(raw, validate=True)
        if len(payload) > MAX_IMAGE_BYTES:
            raise HTTPException(413, f"image_base64 exceeds the {MAX_IMAGE_BYTES}-byte limit")
        image = Image.open(io.BytesIO(payload))
        image.load()
        return image.convert("RGB")
    except HTTPException:
        raise
    except (ValueError, UnidentifiedImageError, OSError, Image.DecompressionBombError) as error:
        raise HTTPException(400, "image_base64 must be a valid, safe base64-encoded image") from error


def make_contents(messages: list[Message]) -> list[dict[str, Any]]:
    result: list[dict[str, Any]] = []
    image_count = 0
    for message in messages:
        if isinstance(message.content, str):
            if len(message.content) > 32_000:
                raise HTTPException(400, "A text message may not exceed 32,000 characters")
            content: str | list[dict[str, Any]] = message.content
        else:
            content = []
            for item in message.content:
                if item.type == "text":
                    content.append({"type": "text", "text": item.text or ""})
                elif item.type == "image":
                    image_count += 1
                    if image_count > 4:
                        raise HTTPException(400, "A request may contain at most four images")
                    if not item.image_base64:
                        raise HTTPException(400, "image content requires image_base64")
                    content.append({"type": "image", "image": decode_image(item.image_base64)})
                else:
                    raise HTTPException(400, f"Unsupported content type: {item.type!r}. Only text and image are supported.")
        result.append({"role": message.role, "content": content})
    return result


@asynccontextmanager
async def lifespan(_app: FastAPI):
    global MODEL, PROCESSOR
    PROCESSOR = AutoProcessor.from_pretrained(str(MODEL_DIR), trust_remote_code=True)
    MODEL = AutoModelForCausalLM.from_pretrained(
        str(MODEL_DIR),
        trust_remote_code=True,
        dtype=torch.bfloat16,
        low_cpu_mem_usage=True,
    ).eval()
    yield
    MODEL = None
    PROCESSOR = None


app = FastAPI(
    title="Phillnet Mini Text-Vision",
    version="1.1.0",
    description="Text generation and still-image understanding only. SDXL, image generation, video generation, audio, tools, and agent runtimes are disabled.",
    lifespan=lifespan,
)

if CORS_ORIGINS:
    app.add_middleware(
        CORSMiddleware,
        allow_origins=CORS_ORIGINS,
        allow_credentials=False,
        allow_methods=["GET", "POST"],
        allow_headers=["Authorization", "Content-Type", "X-API-Key"],
        max_age=600,
    )


@app.middleware("http")
async def enforce_request_limit(request: Request, call_next: Any) -> Any:
    content_length = request.headers.get("content-length")
    if content_length and int(content_length) > MAX_REQUEST_BYTES:
        return JSONResponse(status_code=413, content={"detail": "Request body exceeds configured size limit"})
    return await call_next(request)


@app.get("/health")
def health() -> dict[str, Any]:
    ready = MODEL is not None and PROCESSOR is not None
    return {
        "status": "ok" if ready else "loading",
        "ready": ready,
        "service": "phillnet-mini-text-vision",
        "version": app.version,
        "capabilities": ["text-generation", "image-understanding"],
        "disabled": ["image-generation", "video-generation", "audio", "tools", "agents"],
    }


@app.get("/ready")
def ready() -> dict[str, bool]:
    if MODEL is None or PROCESSOR is None:
        raise HTTPException(503, "Model is still loading")
    return {"ready": True}


@app.post("/v1/chat/completions", dependencies=[Depends(require_api_key)])
def chat_completions(request: ChatRequest) -> dict[str, Any]:
    if MODEL is None or PROCESSOR is None:
        raise HTTPException(503, "Model is still loading")
    content = make_contents(request.messages)
    encoded = PROCESSOR.apply_chat_template(
        content,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    )
    device = next(MODEL.parameters()).device
    encoded = {key: value.to(device) if torch.is_tensor(value) else value for key, value in dict(encoded).items()}
    prompt_length = int(encoded["input_ids"].shape[1])
    generation_kwargs: dict[str, Any] = {
        "max_new_tokens": request.max_tokens,
        "do_sample": request.temperature > 0.0,
        "use_cache": True,
        "reasoning_effort": request.reasoning_effort,
    }
    if request.temperature > 0.0:
        generation_kwargs["temperature"] = request.temperature
    started = time.perf_counter()
    # A single local model instance should perform one generation at a time to
    # prevent concurrent high-context calls from overcommitting model memory.
    with GENERATION_LOCK, torch.inference_mode():
        output = MODEL.generate(**encoded, **generation_kwargs)
    completion_ids = output[0, prompt_length:].detach().cpu()
    text = PROCESSOR.tokenizer.decode(completion_ids, skip_special_tokens=True)
    completion_tokens = int(completion_ids.numel())
    return {
        "id": f"chatcmpl-{uuid.uuid4().hex}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": "phillnet-mini-text-vision",
        "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}],
        "usage": {"prompt_tokens": prompt_length, "completion_tokens": completion_tokens, "total_tokens": prompt_length + completion_tokens},
        "elapsed_seconds": round(time.perf_counter() - started, 3),
    }