"""OCR extraction resource for CMDOP LLM service."""

from __future__ import annotations

from typing import TYPE_CHECKING, Literal, Optional, Union

import httpx
from openai._compat import cached_property
from openai._resource import AsyncAPIResource, SyncAPIResource
from openai._types import NOT_GIVEN, Body, Headers, NotGiven, Query

from ..types.ocr import OCRResponse

if TYPE_CHECKING:
    from .._client import AsyncCmdopLLM, CmdopLLM

__all__ = ["OCR", "AsyncOCR"]


class OCR(SyncAPIResource):
    """OCR text extraction resource.

    Provides text extraction from images using vision models.

    Example:
        ```python
        from cmdop_llm import CmdopLLM

        client = CmdopLLM()
        result = client.ocr.extract(
            image_url="https://example.com/document.png"
        )
        print(result.text)
        ```
    """

    _client: CmdopLLM

    @cached_property
    def with_raw_response(self) -> OCRWithRawResponse:
        """Access raw HTTP response data."""
        return OCRWithRawResponse(self)

    def extract(
        self,
        *,
        image: Optional[str] = None,
        image_url: Optional[str] = None,
        model: Optional[str] = None,
        mode: Literal["tiny", "small", "base", "gundam"] = "base",
        fetch_image: bool = True,
        extra_headers: Headers | None = None,
        extra_query: Query | None = None,
        extra_body: Body | None = None,
        timeout: Union[float, httpx.Timeout, None, NotGiven] = NOT_GIVEN,
    ) -> OCRResponse:
        """Extract text from an image using OCR.

        Args:
            image: Base64-encoded image data.
            image_url: URL of the image to extract text from.
            model: Model to use. Defaults to server default.
            mode: OCR quality mode ("tiny", "small", "base", "gundam").
            fetch_image: Whether server should fetch image from URL.
            extra_headers: Additional headers to send.
            extra_query: Additional query parameters.
            extra_body: Additional body parameters.
            timeout: Request timeout.

        Returns:
            OCRResponse with extracted text.

        Raises:
            ValueError: If neither image nor image_url is provided.
            openai.APIError: If the request fails.
        """
        if not image and not image_url:
            raise ValueError("Either 'image' or 'image_url' must be provided")

        body = {
            "image": image,
            "image_url": image_url,
            "model": model,
            "mode": mode,
            "fetch_image": fetch_image,
        }
        # Remove None values
        body = {k: v for k, v in body.items() if v is not None}

        return self._post(
            "/vision/ocr",
            body=body,
            options={
                "extra_headers": extra_headers,
                "extra_query": extra_query,
                "extra_body": extra_body,
                "timeout": timeout,
            },
            cast_to=OCRResponse,
        )


class AsyncOCR(AsyncAPIResource):
    """Async OCR text extraction resource."""

    _client: AsyncCmdopLLM

    @cached_property
    def with_raw_response(self) -> AsyncOCRWithRawResponse:
        """Access raw HTTP response data."""
        return AsyncOCRWithRawResponse(self)

    async def extract(
        self,
        *,
        image: Optional[str] = None,
        image_url: Optional[str] = None,
        model: Optional[str] = None,
        mode: Literal["tiny", "small", "base", "gundam"] = "base",
        fetch_image: bool = True,
        extra_headers: Headers | None = None,
        extra_query: Query | None = None,
        extra_body: Body | None = None,
        timeout: Union[float, httpx.Timeout, None, NotGiven] = NOT_GIVEN,
    ) -> OCRResponse:
        """Extract text from an image using OCR (async).

        See OCR.extract for full documentation.
        """
        if not image and not image_url:
            raise ValueError("Either 'image' or 'image_url' must be provided")

        body = {
            "image": image,
            "image_url": image_url,
            "model": model,
            "mode": mode,
            "fetch_image": fetch_image,
        }
        body = {k: v for k, v in body.items() if v is not None}

        return await self._post(
            "/vision/ocr",
            body=body,
            options={
                "extra_headers": extra_headers,
                "extra_query": extra_query,
                "extra_body": extra_body,
                "timeout": timeout,
            },
            cast_to=OCRResponse,
        )


class OCRWithRawResponse:
    """Raw response wrapper for OCR resource."""

    def __init__(self, ocr: OCR) -> None:
        self._ocr = ocr

    def extract(self, **kwargs) -> httpx.Response:  # type: ignore
        """Get raw response from extract."""
        return self._ocr._client.with_raw_response._post(  # type: ignore
            "/vision/ocr",
            body=kwargs,
        )


class AsyncOCRWithRawResponse:
    """Raw response wrapper for AsyncOCR resource."""

    def __init__(self, ocr: AsyncOCR) -> None:
        self._ocr = ocr

    async def extract(self, **kwargs) -> httpx.Response:  # type: ignore
        """Get raw response from extract."""
        return await self._ocr._client.with_raw_response._post(  # type: ignore
            "/vision/ocr",
            body=kwargs,
        )
