import base64
import io
import json
from copy import deepcopy
from importlib.resources import files
from urllib.parse import urlparse

from PIL import Image as PILImage

from .. import RELAIData


def _get_relai_client_cls():
    try:
        from relai_core import RELAIClient
    except ModuleNotFoundError as exc:
        raise ModuleNotFoundError(
            "Image uploads via `upload_to_run_id` require a compatible `relai-core` installation."
        ) from exc
    return RELAIClient


def is_valid_url(url: str) -> bool:
    """Check if the URL is valid.

    Args:
        url (str): URL to be validated.

    Returns:
        bool: True if the URL is valid, False otherwise.
    """
    parsed_url = urlparse(url)
    return all([parsed_url.scheme, parsed_url.netloc])


class Image(RELAIData):
    """Data model for an image component. Either `image_url` or `image` must be provided.

    Args:
        image_url (str, optional): URL of the image. Defaults to None.
        image (PILImage.Image, optional): Image object. Defaults to None.
        upload_to_run_id (str, optional): Run ID to which the image should be uploaded.
            Setting this parameter will first upload the image to the specified run ID and then use the image URL.
            This significantly reduces the size of the JSON payload. Defaults to None.
        caption (str, optional): Caption for the image. Defaults to None.
        column_span (int, optional): Number of columns the component should span. Defaults to 12.
    """

    schema_file: str = files("relai_visualization").joinpath("schema/components/image.schema.json")

    def __init__(
        self,
        image_url: str = None,
        image: PILImage.Image = None,
        upload_to_run_id: str = None,
        caption: str = None,
        column_span: int = 12,
    ) -> None:
        if (image_url and image) or (not image_url and not image):
            raise ValueError("Exactly one of `image_url` or `image` must be provided.")

        if image_url and upload_to_run_id:
            raise ValueError("Cannot provide both `image_url` and `upload_to_run_id`.")

        image_data = None
        if image:
            # Upload image and retrieve URL
            if upload_to_run_id:
                relai_client = _get_relai_client_cls()()
                image_url = relai_client.upload_image(image, run_id=upload_to_run_id)
            else:
                buffer = io.BytesIO()
                image.save(buffer, format="png")
                image_data = base64.b64encode(buffer.getvalue()).decode("utf-8")

        super().__init__(
            componentName="image",
            columnSpan=column_span,
            src=image_url,
            data=image_data,
            caption=caption,
        )

    def custom_validation(self) -> None:
        try:
            if not is_valid_url(self.data["src"]):
                raise ValueError("Invalid URL")
        except KeyError:
            # Skip validation if src is not provided
            pass

    def __str__(self) -> str:
        output_dict = deepcopy(self.data)
        try:
            output_dict["data"] = f"Base64 encoded image: '{output_dict['data'][:10]}...'"
        except KeyError:
            pass
        return json.dumps(output_dict)
