# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import asyncio
import base64
import binascii
import logging
import os
from collections import OrderedDict
from io import BytesIO
from typing import Any, Dict, Final, List
from urllib.parse import urlparse

import httpx
from PIL import Image

import dynamo.nixl_connect as nixl_connect
from dynamo.common.utils import nvtx_utils as _nvtx
from dynamo.common.utils.media_nixl import read_decoded_media_via_nixl
from dynamo.common.utils.runtime import run_async

from .http_client import get_http_client

logger = logging.getLogger(__name__)

# Constants for multimodal data variants
URL_VARIANT_KEY: Final = "Url"
DECODED_VARIANT_KEY: Final = "Decoded"


class ImageLoader:
    CACHE_SIZE_MAXIMUM = int(os.environ.get("DYN_MM_IMAGE_CACHE_SIZE", "8"))

    def __init__(
        self,
        cache_size: int = CACHE_SIZE_MAXIMUM,
        http_timeout: float = 30.0,
        enable_frontend_decoding: bool = False,
    ):
        """
        Initialize the ImageLoader with caching, HTTP settings, and optional NIXL config for
        receiving frontend decoding.

        Args:
            cache_size: Maximum number of images to store in the in-memory LRU cache.
                Defaults to CACHE_SIZE_MAXIMUM.
            http_timeout: Timeout in seconds for HTTP requests when fetching remote images.
                Defaults to 30.0 seconds.
            enable_frontend_decoding: If True, enables NIXL RDMA for transferring
                decoded images directly from frontend memory, bypassing standard
                network transport. Defaults to False.
        """
        self._http_timeout = http_timeout
        self._cache_size = cache_size
        self._image_cache: OrderedDict[str, Image.Image] = OrderedDict()
        self._inflight: dict[str, asyncio.Task[Image.Image]] = {}
        self._enable_frontend_decoding = enable_frontend_decoding
        # Lazy-init NIXL connector only when frontend decoding is enabled
        self._nixl_connector = None
        if self._enable_frontend_decoding:
            self._nixl_connector = nixl_connect.Connector()
            run_async(
                self._nixl_connector.initialize
            )  # Synchronously wait for async init

    @staticmethod
    def _open_image_sync(image_data: BytesIO) -> Image.Image:
        """Open, validate, and decode an image from raw bytes. Runs in a thread."""
        image = Image.open(image_data, formats=["JPEG", "PNG", "WEBP"])
        if image.format not in ("JPEG", "PNG", "WEBP"):
            raise ValueError(f"Unsupported image format: {image.format}")
        # Image.open() is lazy — convert() forces the actual pixel decode
        return image.convert("RGB")

    @staticmethod
    async def _open_image(image_data: BytesIO) -> Image.Image:
        """Open and validate an image from raw bytes, converting to RGB."""
        with _nvtx.annotate("mm:img:pil_open_convert", color="lime"):
            return await asyncio.to_thread(ImageLoader._open_image_sync, image_data)

    def _cache_put(self, key: str, image: Image.Image) -> None:
        """Insert into cache if not already present. Sync — no awaits."""
        if key not in self._image_cache:
            if len(self._image_cache) >= self._cache_size:
                self._image_cache.popitem(last=False)
            self._image_cache[key] = image

    async def _fetch_and_process(self, image_url: str) -> Image.Image:
        """Fetch image via HTTP(S), decode with PIL, return RGB Image.

        All exception normalization happens here so shared callers
        see identical error types.
        """
        try:
            with _nvtx.annotate("mm:img:http_fetch", color="lime"):
                http_client = get_http_client(self._http_timeout)
                response = await http_client.get(image_url)
                response.raise_for_status()
                if not response.content:
                    raise ValueError("Empty response content from image URL")
                image_data = BytesIO(response.content)

            return await self._open_image(image_data)

        except httpx.HTTPStatusError as e:
            logger.error(f"HTTP {e.response.status_code} loading image: '{image_url}'")
            raise
        except httpx.TimeoutException as e:
            logger.error(
                f"{type(e).__name__} loading image: '{image_url}' "
                f"(timeout={self._http_timeout}s)"
            )
            raise ValueError(f"Timeout loading image: '{image_url}'") from e
        except httpx.HTTPError as e:
            logger.error(f"{type(e).__name__} loading image: '{image_url}': {e}")
            raise
        except Exception as e:
            logger.error(f"{type(e).__name__} loading image: '{image_url}': {e}")
            raise ValueError(f"Failed to load image: '{image_url}': {e}") from e

    async def _fetch_and_cache(self, key: str, image_url: str) -> Image.Image:
        """Shared task: fetch, cache, then remove from _inflight."""
        try:
            image = await self._fetch_and_process(image_url)
            self._cache_put(key, image)
            return image
        finally:
            self._inflight.pop(key, None)

    @_nvtx.annotate("mm:img:load_image", color="lime")
    async def load_image(self, image_url: str) -> Image.Image:
        parsed_url = urlparse(image_url)

        if parsed_url.scheme in ("http", "https"):
            key = image_url.lower()

            # Check cache (sync — no await, no interleaving possible)
            if key in self._image_cache:
                logger.debug(f"Image found in cache for URL: {image_url}")
                self._image_cache.move_to_end(key)
                return self._image_cache[key]

            # Join existing in-flight task, or start a new one
            if key not in self._inflight:
                task = asyncio.create_task(self._fetch_and_cache(key, image_url))
                # Suppress "exception was never retrieved" if all waiters cancel
                task.add_done_callback(
                    lambda t: t.exception() if not t.cancelled() else None
                )
                self._inflight[key] = task

            # shield so cancelling THIS caller doesn't cancel the shared task
            return await asyncio.shield(self._inflight[key])

        try:
            if parsed_url.scheme == "data":
                with _nvtx.annotate("mm:img:base64_decode", color="lime"):
                    if not parsed_url.path.startswith("image/"):
                        raise ValueError("Data URL must be an image type")

                    media_type, data = parsed_url.path.split(",", 1)
                    if ";base64" not in media_type:
                        raise ValueError("Data URL must be base64 encoded")

                    try:
                        image_bytes = base64.b64decode(data, validate=True)
                    except binascii.Error as e:
                        raise ValueError(f"Invalid base64 encoding: {e}") from e
                    image_data = BytesIO(image_bytes)

            elif parsed_url.scheme in ("", "file"):
                path = image_url if parsed_url.scheme == "" else parsed_url.path

                def _read_local_file(p: str) -> bytes:
                    with open(p, "rb") as f:
                        return f.read()

                image_bytes = await asyncio.to_thread(_read_local_file, path)
                image_data = BytesIO(image_bytes)
            else:
                raise ValueError(f"Invalid image source scheme: {parsed_url.scheme}")

            return await self._open_image(image_data)

        except Exception as e:
            logger.error(f"{type(e).__name__} loading image: '{image_url}': {e}")
            raise ValueError(f"Failed to load image: '{image_url}': {e}") from e

    async def load_image_batch(
        self,
        image_mm_items: List[Dict[str, Any]],
    ) -> List[Any]:
        """
        Load a batch of images from multimodal data items.

        Supports two paths:
        1. Url variant: Download and decode image from URL (default)
        2. Decoded variant: Read pre-decoded image via NIXL RDMA (requires enable_frontend_decoding=True)

        Args:
            image_mm_items: List of multimodal data items for images

        Returns:
            List of loaded image data

        Raises:
            Exception: If any image fails to load
            ValueError: If enable_frontend_decoding=True but nixl_connector is None
        """
        image_futures = []

        for item in image_mm_items:
            if isinstance(item, dict) and URL_VARIANT_KEY in item:
                # URL path: download and decode in Python backend
                url = item[URL_VARIANT_KEY]
                image_futures.append(self.load_image(url))
                logger.debug(f"Preparing to load image from URL: {url[:80]}...")
            elif isinstance(item, dict) and DECODED_VARIANT_KEY in item:
                if self._enable_frontend_decoding:
                    metadata = item[DECODED_VARIANT_KEY]
                    if self._nixl_connector is None:
                        raise RuntimeError("NIXL connector is not initialized")
                    image_futures.append(
                        read_decoded_media_via_nixl(self._nixl_connector, metadata)
                    )
                else:
                    logger.error(
                        "Received Decoded multimodal data but enable_frontend_decoding=False. "
                        "Set enable_frontend_decoding=True to enable NIXL RDMA image transfer."
                    )
                    raise ValueError("Could not load decoded media from frontend")

        # Process images in parallel
        results = await asyncio.gather(*image_futures, return_exceptions=True)
        loaded_images = []
        collective_exceptions = ""
        for media_item, result in zip(image_mm_items, results):
            if isinstance(result, Exception):
                source = media_item.get(URL_VARIANT_KEY, "decoded")
                logger.error(f"Failed to load image from {source[:80]}...: {result}")
                collective_exceptions += (
                    f"Failed to load image from {source[:80]}...: {result}\n"
                )
                continue
            loaded_images.append(result)

        if collective_exceptions:
            raise Exception(collective_exceptions)

        return loaded_images
