"""RapidOCR client wrappers."""

from __future__ import annotations

from typing import Unpack

from pydantic import BaseModel, Field

from .client import AsyncClient, ClientKwargs
from .common import patch_response_class


class OCRRequest(BaseModel):
  """RapidOCR request payload."""

  image: str = Field(..., description='image URL or data URI input')
  text_score: float = Field(0.5, ge=0, le=1)
  use_det: bool = True
  use_cls: bool = True
  use_rec: bool = True
  return_word_box: bool = False
  return_single_char_box: bool = False


class OCRTextResult(BaseModel):
  """Single OCR text result."""

  box: list[list[float]] = Field(description='text box coordinates')
  text: str = Field(description='recognized text')
  score: float = Field(description='recognition confidence')


class OCRResponse(BaseModel):
  """RapidOCR response payload."""

  results: list[OCRTextResult] = Field(default_factory=list)
  elapse: float
  elapse_list: list[float]


class AsyncRapidOCRClient(AsyncClient):
  """Async client for the RapidOCR service."""

  def __init__(self, base_url: str, **kwargs: Unpack[ClientKwargs]):
    """Initialize the client with the provider base URL."""
    super().__init__(base_url=base_url.rstrip('/'), **kwargs)

  @patch_response_class(model=OCRResponse)
  async def ocr(
    self,
    *,
    image: str,
    text_score: float = 0.5,
    use_det: bool = True,
    use_cls: bool = True,
    use_rec: bool = True,
    return_word_box: bool = False,
    return_single_char_box: bool = False,
    **kwargs,
  ):
    """Run OCR with flattened request parameters."""
    return await self.post(
      '/ocr',
      json={
        'image': image,
        'text_score': text_score,
        'use_det': use_det,
        'use_cls': use_cls,
        'use_rec': use_rec,
        'return_word_box': return_word_box,
        'return_single_char_box': return_single_char_box,
      },
      **kwargs,
    )
