"""
Azure Document Intelligence OCR transformation implementation.

Azure Document Intelligence (formerly Form Recognizer) provides advanced document analysis capabilities.
This implementation transforms between Mistral OCR format and Azure Document Intelligence API v4.0.

Note: Azure Document Intelligence API is async - POST returns 202 Accepted with Operation-Location header.
The operation location must be polled until the analysis completes.
"""

import asyncio
import re
import time
from typing import Any, Dict, Optional
from urllib.parse import quote

import httpx

from litellm._logging import verbose_logger
from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin
from litellm.constants import (
    AZURE_DOCUMENT_INTELLIGENCE_API_VERSION,
    AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI,
    AZURE_OPERATION_POLLING_TIMEOUT,
)
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
from litellm.llms.base_llm.ocr.transformation import (
    BaseOCRConfig,
    DocumentType,
    OCRPage,
    OCRPageDimensions,
    OCRRequestData,
    OCRResponse,
    OCRUsageInfo,
)
from litellm.secret_managers.main import get_secret_str


class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
    """
    Azure Document Intelligence OCR transformation configuration.

    Supports Azure Document Intelligence v4.0 (2024-11-30) API.
    Model route: azure_ai/doc-intelligence/<model>

    Supported models:
    - prebuilt-layout: Extracts text with markdown, tables, and structure (closest to Mistral OCR)
    - prebuilt-read: Basic text extraction optimized for reading
    - prebuilt-document: General document analysis

    Reference: https://learn.microsoft.com/en-us/azure/ai-services/document-intelligence/
    """

    def __init__(self) -> None:
        super().__init__()

    def get_supported_ocr_params(self, model: str) -> list:
        """
        Get supported OCR parameters for Azure Document Intelligence.

        Azure DI exposes a `pages` query parameter on the analyze endpoint
        (1-based, e.g. "1-3,5,7-9"). To keep the public request shape
        aligned with Mistral OCR, callers pass `pages` using Mistral
        semantics — a list of 0-based integers — or a pre-formatted
        Azure-style string. Other Mistral-specific params (e.g.
        `include_image_base64`) are not supported by Azure DI and are
        ignored during transformation.
        """
        return ["pages"]

    def map_ocr_params(
        self,
        non_default_params: dict,
        optional_params: dict,
        model: str,
    ) -> dict:
        """
        Map OCR params to Azure DI format.

        Translates Mistral-style `pages` (list[int], 0-based) into Azure's
        `pages` query string (1-based, e.g. "1,2,3" or "1-3,5"). A raw
        string that already matches Azure's format is passed through
        unchanged.
        """
        pages = non_default_params.get("pages")
        if pages is None:
            return optional_params

        normalized = self._normalize_pages_param(pages)
        if normalized:
            optional_params["pages"] = normalized
        return optional_params

    @staticmethod
    def _normalize_pages_param(pages: Any) -> str:
        """
        Convert a caller-provided `pages` value to Azure DI's query-string
        form. Azure expects 1-based page numbers, grammar: `^(\\d+(-\\d+)?)(,\\s*(\\d+(-\\d+)?))*$`.

        Accepted inputs:
          - list[int]: Mistral-style 0-based indices. Converted to 1-based
            and joined (e.g. [0,1,2] -> "1,2,3").
          - list[str]: tokens like "1" or "3-5". Validated, joined as-is
            (treated as Azure-native, i.e. 1-based).
          - str: already in Azure format. Validated and whitespace-stripped.
        """
        pages_pattern = re.compile(r"^\s*\d+(-\d+)?(\s*,\s*\d+(-\d+)?)*\s*$")

        if isinstance(pages, str):
            if not pages_pattern.match(pages):
                raise ValueError(
                    f"Invalid `pages` string for Azure Document Intelligence: "
                    f"{pages!r}. Expected format like '1-3,5,7-9'."
                )
            return pages.replace(" ", "")

        if isinstance(pages, list):
            if len(pages) == 0:
                return ""
            if any(isinstance(p, bool) for p in pages):
                raise ValueError("`pages` must be integers, not booleans")
            if all(isinstance(p, int) for p in pages):
                if any(p < 0 for p in pages):
                    raise ValueError(
                        "`pages` integers must be >= 0 (Mistral 0-based indices)"
                    )
                # Mistral 0-based -> Azure 1-based.
                return ",".join(str(p + 1) for p in sorted(set(pages)))
            if all(isinstance(p, str) for p in pages):
                joined = ",".join(p.strip() for p in pages)
                if not pages_pattern.match(joined):
                    raise ValueError(
                        f"Invalid `pages` list for Azure Document Intelligence: "
                        f"{pages!r}. Expected tokens like '1' or '3-5'."
                    )
                return joined

        raise ValueError(
            "`pages` must be a list[int] (0-based, Mistral-style) or a "
            "string like '1-3,5,7-9'."
        )

    def validate_environment(
        self,
        headers: Dict,
        model: str,
        api_key: Optional[str] = None,
        api_base: Optional[str] = None,
        litellm_params: Optional[dict] = None,
        **kwargs,
    ) -> Dict:
        """
        Validate environment and return headers for Azure Document Intelligence.

        Authentication uses Ocp-Apim-Subscription-Key header.
        """
        # Get API key from environment if not provided
        if api_key is None:
            api_key = get_secret_str("AZURE_DOCUMENT_INTELLIGENCE_API_KEY")

        if api_key is None:
            raise ValueError(
                "Missing Azure Document Intelligence API Key - Set AZURE_DOCUMENT_INTELLIGENCE_API_KEY environment variable or pass api_key parameter"
            )

        # Validate API base/endpoint is provided
        if api_base is None:
            api_base = get_secret_str("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT")

        if api_base is None:
            raise ValueError(
                "Missing Azure Document Intelligence Endpoint - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT environment variable or pass api_base parameter"
            )

        headers = {
            "Ocp-Apim-Subscription-Key": api_key,
            "Content-Type": "application/json",
            **headers,
        }

        return headers

    def get_complete_url(
        self,
        api_base: Optional[str],
        model: str,
        optional_params: dict,
        litellm_params: Optional[dict] = None,
        **kwargs,
    ) -> str:
        """
        Get complete URL for Azure Document Intelligence endpoint.

        Format: {endpoint}/documentintelligence/documentModels/{modelId}:analyze?api-version=2024-11-30

        Note: API version 2024-11-30 uses /documentintelligence/ path (not /formrecognizer/)

        Args:
            api_base: Azure Document Intelligence endpoint (e.g., https://your-resource.cognitiveservices.azure.com)
            model: Model ID (e.g., "prebuilt-layout", "prebuilt-read")
            optional_params: Optional parameters

        Returns: Complete URL for Azure DI analyze endpoint
        """
        if api_base is None:
            api_base = get_secret_str("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT")

        if api_base is None:
            raise ValueError(
                "Missing Azure Document Intelligence Endpoint - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT environment variable or pass api_base parameter"
            )

        # Ensure no trailing slash
        api_base = api_base.rstrip("/")

        # Extract model ID from full model path if needed
        # Model can be "prebuilt-layout" or "azure_ai/doc-intelligence/prebuilt-layout"
        model_id = model
        if "/" in model:
            # Extract the last part after the last slash
            model_id = model.split("/")[-1]
        encoded_model_id = encode_url_path_segment(model_id, field_name="model_id")

        # Azure Document Intelligence analyze endpoint
        # Note: API version 2024-11-30+ uses /documentintelligence/ (not /formrecognizer/)
        url = (
            f"{api_base}/documentintelligence/documentModels/{encoded_model_id}:analyze"
            f"?api-version={AZURE_DOCUMENT_INTELLIGENCE_API_VERSION}"
        )

        # Azure DI accepts `pages` as a query param (1-based, e.g. "1-3,5").
        # `optional_params` has already been normalized in `map_ocr_params`.
        pages = optional_params.get("pages") if optional_params else None
        if pages:
            url += f"&pages={quote(str(pages), safe=',-')}"

        return url

    def _extract_base64_from_data_uri(self, data_uri: str) -> str:
        """
        Extract base64 content from a data URI.

        Args:
            data_uri: Data URI like "data:application/pdf;base64,..."

        Returns:
            Base64 string without the data URI prefix
        """
        # Match pattern: data:[<mediatype>][;base64],<data>
        match = re.match(r"data:([^;]+)(?:;base64)?,(.+)", data_uri)
        if match:
            return match.group(2)
        return data_uri

    def transform_ocr_request(
        self,
        model: str,
        document: DocumentType,
        optional_params: dict,
        headers: dict,
        **kwargs,
    ) -> OCRRequestData:
        """
        Transform OCR request to Azure Document Intelligence format.

        Mistral OCR format:
        {
            "document": {
                "type": "document_url",
                "document_url": "https://example.com/doc.pdf"
            }
        }

        Azure DI format:
        {
            "urlSource": "https://example.com/doc.pdf"
        }
        OR
        {
            "base64Source": "base64_encoded_content"
        }

        Args:
            model: Model name
            document: Document dict from user (Mistral format)
            optional_params: Already mapped optional parameters
            headers: Request headers

        Returns:
            OCRRequestData with JSON data
        """
        verbose_logger.debug(
            f"Azure Document Intelligence transform_ocr_request - model: {model}"
        )

        if not isinstance(document, dict):
            raise ValueError(f"Expected document dict, got {type(document)}")

        # Extract document URL from Mistral format
        doc_type = document.get("type")
        document_url = None

        if doc_type == "document_url":
            document_url = document.get("document_url", "")
        elif doc_type == "image_url":
            document_url = document.get("image_url", "")
        else:
            raise ValueError(
                f"Invalid document type: {doc_type}. Must be 'document_url' or 'image_url'"
            )

        if not document_url:
            raise ValueError("Document URL is required")

        # Build Azure DI request
        data: Dict[str, Any] = {}

        # Check if it's a data URI (base64)
        if document_url.startswith("data:"):
            # Extract base64 content
            base64_content = self._extract_base64_from_data_uri(document_url)
            data["base64Source"] = base64_content
            verbose_logger.debug("Using base64Source for Azure Document Intelligence")
        else:
            # Regular URL
            data["urlSource"] = document_url
            verbose_logger.debug("Using urlSource for Azure Document Intelligence")

        # Azure DI: `pages` is a query param (wired in get_complete_url),
        # not a body field. Other Mistral-specific params (e.g.
        # include_image_base64, image_limit) are unsupported and ignored.

        return OCRRequestData(data=data, files=None)

    def _extract_page_markdown(self, page_data: Dict[str, Any]) -> str:
        """
        Extract text from Azure DI page and format as markdown.

        Azure DI provides text in 'lines' array. We concatenate them with newlines.

        Args:
            page_data: Azure DI page object

        Returns:
            Markdown-formatted text
        """
        lines = page_data.get("lines", [])
        if not lines:
            return ""

        # Extract text content from each line
        text_lines = [line.get("content", "") for line in lines]

        # Join with newlines to preserve structure
        return "\n".join(text_lines)

    def _convert_dimensions(
        self, width: float, height: float, unit: str
    ) -> OCRPageDimensions:
        """
        Convert Azure DI dimensions to pixels.

        Azure DI provides dimensions in inches. We convert to pixels using configured DPI.

        Args:
            width: Width in specified unit
            height: Height in specified unit
            unit: Unit of measurement (e.g., "inch")

        Returns:
            OCRPageDimensions with pixel values
        """
        # Convert to pixels using configured DPI
        dpi = AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI
        if unit == "inch":
            width_px = int(width * dpi)
            height_px = int(height * dpi)
        else:
            # If unit is not inches, assume it's already in pixels
            width_px = int(width)
            height_px = int(height)

        return OCRPageDimensions(width=width_px, height=height_px, dpi=dpi)

    @staticmethod
    def _check_timeout(start_time: float, timeout_secs: int) -> None:
        """
        Check if operation has timed out.

        Args:
            start_time: Start time of the operation
            timeout_secs: Timeout duration in seconds

        Raises:
            TimeoutError: If operation has exceeded timeout
        """
        if time.time() - start_time > timeout_secs:
            raise TimeoutError(
                f"Azure Document Intelligence operation polling timed out after {timeout_secs} seconds"
            )

    @staticmethod
    def _get_retry_after(response: httpx.Response) -> int:
        """
        Get retry-after duration from response headers.

        Args:
            response: HTTP response

        Returns:
            Retry-after duration in seconds (default: 2)
        """
        retry_after = int(response.headers.get("retry-after", "2"))
        verbose_logger.debug(f"Retry polling after: {retry_after} seconds")
        return retry_after

    @staticmethod
    def _check_operation_status(response: httpx.Response) -> str:
        """
        Check Azure DI operation status from response.

        Args:
            response: HTTP response from operation endpoint

        Returns:
            Operation status string

        Raises:
            ValueError: If operation failed or status is unknown
        """
        try:
            result = response.json()
            status = result.get("status")

            verbose_logger.debug(f"Azure DI operation status: {status}")

            if status == "succeeded":
                return "succeeded"
            elif status == "failed":
                error_msg = result.get("error", {}).get("message", "Unknown error")
                raise ValueError(
                    f"Azure Document Intelligence analysis failed: {error_msg}"
                )
            elif status in ["running", "notStarted"]:
                return "running"
            else:
                raise ValueError(f"Unknown operation status: {status}")

        except Exception as e:
            if "succeeded" in str(e) or "failed" in str(e):
                raise
            # If we can't parse JSON, something went wrong
            raise ValueError(f"Failed to parse Azure DI operation response: {e}")

    def _poll_operation_sync(
        self,
        operation_url: str,
        headers: Dict[str, str],
        timeout_secs: int,
    ) -> httpx.Response:
        """
        Poll Azure Document Intelligence operation until completion (sync).

        Azure DI POST returns 202 with Operation-Location header.
        We need to poll that URL until status is "succeeded" or "failed".

        Args:
            operation_url: The Operation-Location URL to poll
            headers: Request headers (including auth)
            timeout_secs: Total timeout in seconds

        Returns:
            Final response with completed analysis
        """
        from litellm.llms.custom_httpx.http_handler import _get_httpx_client

        client = _get_httpx_client()
        start_time = time.time()

        verbose_logger.debug(f"Polling Azure DI operation: {operation_url}")

        while True:
            self._check_timeout(start_time=start_time, timeout_secs=timeout_secs)

            # Poll the operation status
            response = client.get(url=operation_url, headers=headers)

            # Check operation status
            status = self._check_operation_status(response=response)

            if status == "succeeded":
                return response
            elif status == "running":
                # Wait before polling again
                retry_after = self._get_retry_after(response=response)
                time.sleep(retry_after)

    async def _poll_operation_async(
        self,
        operation_url: str,
        headers: Dict[str, str],
        timeout_secs: int,
    ) -> httpx.Response:
        """
        Poll Azure Document Intelligence operation until completion (async).

        Args:
            operation_url: The Operation-Location URL to poll
            headers: Request headers (including auth)
            timeout_secs: Total timeout in seconds

        Returns:
            Final response with completed analysis
        """
        import litellm
        from litellm.llms.custom_httpx.http_handler import get_async_httpx_client

        client = get_async_httpx_client(llm_provider=litellm.LlmProviders.AZURE_AI)
        start_time = time.time()

        verbose_logger.debug(f"Polling Azure DI operation (async): {operation_url}")

        while True:
            self._check_timeout(start_time=start_time, timeout_secs=timeout_secs)

            # Poll the operation status
            response = await client.get(url=operation_url, headers=headers)

            # Check operation status
            status = self._check_operation_status(response=response)

            if status == "succeeded":
                return response
            elif status == "running":
                # Wait before polling again
                retry_after = self._get_retry_after(response=response)
                await asyncio.sleep(retry_after)

    def transform_ocr_response(
        self,
        model: str,
        raw_response: httpx.Response,
        logging_obj: Any,
        **kwargs,
    ) -> OCRResponse:
        """
        Transform Azure Document Intelligence response to Mistral OCR format.

        Handles async operation polling: If response is 202 Accepted, polls Operation-Location
        until analysis completes.

        Azure DI response (after polling):
        {
            "status": "succeeded",
            "analyzeResult": {
                "content": "Full document text...",
                "pages": [
                    {
                        "pageNumber": 1,
                        "width": 8.5,
                        "height": 11,
                        "unit": "inch",
                        "lines": [{"content": "text", "boundingBox": [...]}]
                    }
                ]
            }
        }

        Mistral OCR format:
        {
            "pages": [
                {
                    "index": 0,
                    "markdown": "extracted text",
                    "dimensions": {"width": 816, "height": 1056, "dpi": 96}
                }
            ],
            "model": "azure_ai/doc-intelligence/prebuilt-layout",
            "usage_info": {"pages_processed": 1},
            "object": "ocr"
        }

        Args:
            model: Model name
            raw_response: Raw HTTP response from Azure DI (may be 202 Accepted)
            logging_obj: Logging object

        Returns:
            OCRResponse in Mistral format
        """
        try:
            # Check if we got 202 Accepted (async operation started)
            if raw_response.status_code == 202:
                verbose_logger.debug(
                    "Azure DI returned 202 Accepted, polling operation..."
                )

                # Get Operation-Location header
                operation_url = raw_response.headers.get("Operation-Location")
                if not operation_url:
                    raise ValueError(
                        "Azure Document Intelligence returned 202 but no Operation-Location header found"
                    )

                # Reject cross-origin polling URLs — the auth headers
                # below would otherwise leak to whatever URL the upstream
                # (or an attacker-controlled upstream) returns. VERIA-51.
                try:
                    assert_same_origin(operation_url, str(raw_response.request.url))
                except SSRFError as ssrf_err:
                    raise ValueError(
                        f"Azure Document Intelligence: rejected polling URL ({ssrf_err})"
                    )

                # Get headers for polling (need auth)
                poll_headers = {
                    "Ocp-Apim-Subscription-Key": raw_response.request.headers.get(
                        "Ocp-Apim-Subscription-Key", ""
                    )
                }

                # Get timeout from kwargs or use default
                timeout_secs = AZURE_OPERATION_POLLING_TIMEOUT

                # Poll until operation completes
                raw_response = self._poll_operation_sync(
                    operation_url=operation_url,
                    headers=poll_headers,
                    timeout_secs=timeout_secs,
                )

            # Now parse the completed response
            response_json = raw_response.json()

            verbose_logger.debug(
                f"Azure Document Intelligence response status: {response_json.get('status')}"
            )

            # Check if request succeeded
            status = response_json.get("status")
            if status != "succeeded":
                raise ValueError(
                    f"Azure Document Intelligence analysis failed with status: {status}"
                )

            # Extract analyze result
            analyze_result = response_json.get("analyzeResult", {})
            azure_pages = analyze_result.get("pages", [])

            # Transform pages to Mistral format
            mistral_pages = []
            for azure_page in azure_pages:
                page_number = azure_page.get("pageNumber", 1)
                index = page_number - 1  # Convert to 0-based index

                # Extract markdown text
                markdown = self._extract_page_markdown(azure_page)

                # Convert dimensions
                width = azure_page.get("width", 8.5)
                height = azure_page.get("height", 11)
                unit = azure_page.get("unit", "inch")
                dimensions = self._convert_dimensions(
                    width=width, height=height, unit=unit
                )

                # Build OCR page
                ocr_page = OCRPage(
                    index=index, markdown=markdown, dimensions=dimensions
                )
                mistral_pages.append(ocr_page)

            # Build usage info
            usage_info = OCRUsageInfo(
                pages_processed=len(mistral_pages), doc_size_bytes=None
            )

            # Return Mistral OCR response
            return OCRResponse(
                pages=mistral_pages,
                model=model,
                usage_info=usage_info,
                object="ocr",
            )

        except Exception as e:
            verbose_logger.error(
                f"Error parsing Azure Document Intelligence response: {e}"
            )
            raise e

    async def async_transform_ocr_response(
        self,
        model: str,
        raw_response: httpx.Response,
        logging_obj: Any,
        **kwargs,
    ) -> OCRResponse:
        """
        Async transform Azure Document Intelligence response to Mistral OCR format.

        Handles async operation polling: If response is 202 Accepted, polls Operation-Location
        until analysis completes using async polling.

        Args:
            model: Model name
            raw_response: Raw HTTP response from Azure DI (may be 202 Accepted)
            logging_obj: Logging object

        Returns:
            OCRResponse in Mistral format
        """
        try:
            # Check if we got 202 Accepted (async operation started)
            if raw_response.status_code == 202:
                verbose_logger.debug(
                    "Azure DI returned 202 Accepted, polling operation (async)..."
                )

                # Get Operation-Location header
                operation_url = raw_response.headers.get("Operation-Location")
                if not operation_url:
                    raise ValueError(
                        "Azure Document Intelligence returned 202 but no Operation-Location header found"
                    )

                # Reject cross-origin polling URLs (see sync path). VERIA-51.
                try:
                    assert_same_origin(operation_url, str(raw_response.request.url))
                except SSRFError as ssrf_err:
                    raise ValueError(
                        f"Azure Document Intelligence: rejected polling URL ({ssrf_err})"
                    )

                # Get headers for polling (need auth)
                poll_headers = {
                    "Ocp-Apim-Subscription-Key": raw_response.request.headers.get(
                        "Ocp-Apim-Subscription-Key", ""
                    )
                }

                # Get timeout from kwargs or use default
                timeout_secs = AZURE_OPERATION_POLLING_TIMEOUT

                # Poll until operation completes (async)
                raw_response = await self._poll_operation_async(
                    operation_url=operation_url,
                    headers=poll_headers,
                    timeout_secs=timeout_secs,
                )

            # Now parse the completed response
            response_json = raw_response.json()

            verbose_logger.debug(
                f"Azure Document Intelligence response status: {response_json.get('status')}"
            )

            # Check if request succeeded
            status = response_json.get("status")
            if status != "succeeded":
                raise ValueError(
                    f"Azure Document Intelligence analysis failed with status: {status}"
                )

            # Extract analyze result
            analyze_result = response_json.get("analyzeResult", {})
            azure_pages = analyze_result.get("pages", [])

            # Transform pages to Mistral format
            mistral_pages = []
            for azure_page in azure_pages:
                page_number = azure_page.get("pageNumber", 1)
                index = page_number - 1  # Convert to 0-based index

                # Extract markdown text
                markdown = self._extract_page_markdown(azure_page)

                # Convert dimensions
                width = azure_page.get("width", 8.5)
                height = azure_page.get("height", 11)
                unit = azure_page.get("unit", "inch")
                dimensions = self._convert_dimensions(
                    width=width, height=height, unit=unit
                )

                # Build OCR page
                ocr_page = OCRPage(
                    index=index, markdown=markdown, dimensions=dimensions
                )
                mistral_pages.append(ocr_page)

            # Build usage info
            usage_info = OCRUsageInfo(
                pages_processed=len(mistral_pages), doc_size_bytes=None
            )

            # Return Mistral OCR response
            return OCRResponse(
                pages=mistral_pages,
                model=model,
                usage_info=usage_info,
                object="ocr",
            )

        except Exception as e:
            verbose_logger.error(
                f"Error parsing Azure Document Intelligence response (async): {e}"
            )
            raise e
