import time
import uuid
from typing import Any, Dict, List, Literal, Optional, Union

from fastapi import APIRouter, Request
from fastapi.responses import JSONResponse
from pydantic import BaseModel, ConfigDict

from surya.endpoint import service

router = APIRouter()


class _ImageUrl(BaseModel):
    model_config = ConfigDict(extra="allow")
    url: str


class _ContentPart(BaseModel):
    model_config = ConfigDict(extra="allow")
    type: str
    image_url: Optional[_ImageUrl] = None
    text: Optional[str] = None


class _ChatMessage(BaseModel):
    model_config = ConfigDict(extra="allow")
    role: str
    content: Union[str, List[_ContentPart]]


class ChatCompletionRequest(BaseModel):
    # The OpenAI SDK merges `extra_body={...}` into the TOP LEVEL of the request
    # JSON, so the OCR control flags live here rather than nested under a key.
    model_config = ConfigDict(extra="allow")
    model: str = "surya-ocr"
    messages: List[_ChatMessage]
    stream: Optional[bool] = False
    skip_text_detection: bool = False
    recognize_math: bool = False
    skip_table_detection: bool = False
    ocr_with_boxes: bool = True
    mode: Literal["block", "full_page", "table"] = "block"


def _openai_error(message: str, err_type: str, status_code: int) -> JSONResponse:
    return JSONResponse(
        status_code=status_code,
        content={"error": {"message": message, "type": err_type, "code": None}},
    )


def _build_chat_completion(
    content: str, model: str, ocr_json: Any, elapsed: Optional[float]
) -> Dict[str, Any]:
    # Token accounting happens inside vLLM and is not surfaced to this layer, so
    # usage counts are reported as 0 (prompt_tokens intentionally 0).
    return {
        "id": f"chatcmpl-{uuid.uuid4().hex}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": model,
        "choices": [
            {
                "index": 0,
                "finish_reason": "stop",
                "message": {"role": "assistant", "content": content},
            }
        ],
        "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
        "surya": {"ocr_text_json": ocr_json, "elapsed_seconds": elapsed},
    }


@router.post("/v1/chat/completions")
def chat_completions(req: ChatCompletionRequest, request: Request):
    request_id = getattr(request.state, "request_id", "-")
    if req.stream:
        return _openai_error(
            "streaming is not supported by this OCR endpoint",
            "invalid_request_error",
            400,
        )
    try:
        image = service.extract_last_image(req.messages)
    except service.OcrRequestError as exc:
        return _openai_error(str(exc), "invalid_request_error", 400)

    try:
        if req.mode == "table":
            data = service.table_image(image, req.skip_table_detection)
        else:
            skip = req.skip_text_detection or req.mode == "full_page"
            data = service.ocr_via_batcher(
                image,
                skip_text_detection=skip,
                recognize_math=req.recognize_math,
                request_id=request_id,
            )
    except service.OcrRequestError as exc:
        return _openai_error(str(exc), "invalid_request_error", 400)
    except Exception as exc:
        service._log_exception("chat_completions", request, exc, {})
        return _openai_error(str(exc), "internal_error", 500)

    content = data.get("text_lines", "") or ""
    ocr_json = data.get("ocr_text_json") if req.ocr_with_boxes else None
    return _build_chat_completion(content, req.model, ocr_json, data.get("elapsed_seconds"))