"""MarkItDown client wrappers."""

from __future__ import annotations

import base64
from typing import Unpack

from pydantic import BaseModel, Field

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


class ConvertFileResponse(BaseModel):
  """MarkItDown file conversion response."""

  markdown: str = Field(..., description='converted markdown output')
  title: str | None = Field(None, description='document title')
  filename: str = Field(..., description='resolved filename')


class AsyncMarkItDownClient(AsyncClient):
  """Async client for the MarkItDown 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=ConvertFileResponse)
  async def convert_file(
    self,
    *,
    file: tuple[str, bytes, str],
    mime_type: str | None = None,
    **kwargs,
  ):
    """Convert an uploaded file into markdown."""
    data = {'mime_type': mime_type} if mime_type else None
    return await self.post(
      '/convert/file',
      files={'file': file},
      data=data,
      **kwargs,
    )

  @patch_response_class(model=ConvertFileResponse)
  async def convert(
    self,
    *,
    content: str | bytes,
    filename: str | None = None,
    mime_type: str | None = None,
    **kwargs,
  ):
    """Convert a remote URL into markdown."""
    if isinstance(content, bytes):
      content = base64.b64encode(content).decode('utf-8')
    payload: dict[str, str] = {'content': content}
    if filename is not None:
      payload['filename'] = filename
    if mime_type is not None:
      payload['mime_type'] = mime_type
    return await self.post('/convert', json=payload, **kwargs)
