from typing import Any, Callable, Iterator, cast

from pandas import DataFrame, json_normalize

from impresso.api_client.api.images import find_images
from impresso.api_client.api.images import get_image
from impresso.api_client.models.find_images_base_find_response import (
    FindImagesBaseFindResponse,
)
from impresso.api_client.models.find_images_order_by import (
    FindImagesOrderBy,
    FindImagesOrderByLiteral,
)
from impresso.api_client.types import UNSET
from impresso.api_models import Q, BaseFind, Filter, Image
from impresso.data_container import DataContainer, iterate_pages
from impresso.resources.base import DEFAULT_PAGE_SIZE, Resource
from impresso.resources.tools import Embedding
from impresso.structures import AND, OR, DateRange
from impresso.util.error import raise_for_error
from impresso.util.filters import and_or_filter, filters_as_protobuf
from impresso.util.py import get_enum_from_literal
import pandas as pd


# The formatter function needs to be available
def _path_to_image_html(path):
    # Ensure the URL is wrapped in an HTML <img> tag
    return f'<img src="{path}"/>'


class ImageDataFrame(pd.DataFrame):
    @property
    def _constructor(self):
        return ImageDataFrame

    def _repr_html_(self):
        """
        Overrides the default method used by Jupyter/Colab to render DataFrames.
        It forces the display of the styled version.
        """
        # 1. Apply the desired style/formatter
        styler = self.style.format(
            {"previewImage": _path_to_image_html},
        )

        # 2. Return the HTML string generated by the Styler
        return styler.to_html()


def _as_image_df(df: pd.DataFrame) -> ImageDataFrame:
    idf = ImageDataFrame(df)
    cols = idf.columns.tolist()

    if "previewUrl" in cols:
        idf["previewImage"] = idf["previewUrl"]
        cols = idf.columns.tolist()
    # Rearrange columns to make previewImage second
    if "previewImage" in cols:
        cols.remove("previewImage")
        cols.insert(1, "previewImage")
    # Push the URL to the back
    if "previewUrl" in cols:
        cols.remove("previewUrl")
        cols.insert(-1, "previewUrl")
    return cast(ImageDataFrame, idf[cols])


class FindImagesSchema(BaseFind):
    """Schema for the find images response."""

    data: list[Image]


class FindImagesContainer(DataContainer):
    """Response of a find call."""

    def __init__(
        self,
        data: FindImagesBaseFindResponse,
        pydantic_model: type[FindImagesSchema],
        fetch_method: Callable[..., "FindImagesContainer"],
        fetch_method_args: dict[str, Any],
        web_app_search_result_url: str | None = None,
    ):
        super().__init__(data, pydantic_model, web_app_search_result_url)
        self._fetch_method = fetch_method
        self._fetch_method_args = fetch_method_args

    @property
    def df(self) -> DataFrame:
        """Return the data as a pandas dataframe."""
        data = self._data.to_dict()["data"]
        if len(data):
            df = json_normalize(self._data.to_dict()["data"]).set_index("uid")
            return _as_image_df(df)

        return DataFrame()

    def pages(self) -> Iterator["FindImagesContainer"]:
        """Iterate over all pages of results."""
        yield self
        yield from iterate_pages(
            self._fetch_method,
            self._fetch_method_args,
            self.offset,
            self.limit,
            self.total,
        )


class GetImageContainer(DataContainer):

    @property
    def df(self) -> DataFrame:
        """Return the data as a pandas dataframe."""
        data = self._data.to_dict()
        if len(data):
            df = json_normalize(data).set_index("uid")
            return _as_image_df(df)

        return DataFrame()


class ImagesResource(Resource):
    """Search images in Impresso.

    Examples:
        Search for images by keyword:
        >>> results = images.find(term="war")  # doctest: +SKIP
        >>> print(results.df)  # doctest: +SKIP

        Filter images by date range and newspaper:
        >>> from impresso import DateRange
        >>> date_range = DateRange(start="1900-01-01", end="1910-12-31")
        >>> results = images.find(media_id="GDL", date_range=date_range)  # doctest: +SKIP
        >>> print(results.df)  # doctest: +SKIP

        Search for front page images only:
        >>> results = images.find(is_front=True)  # doctest: +SKIP
        >>> print(results.df)  # doctest: +SKIP

        Search images by visual similarity using embeddings:
        >>> embedding = tools.embed_image("path/to/image.jpg", target="image")  # doctest: +SKIP
        >>> similar_images = images.find(embedding=embedding)  # doctest: +SKIP
        >>> print(similar_images.df)  # doctest: +SKIP

        Get a specific image by its ID:
        >>> image_id = "some-image-id"  # Replace with a real ID
        >>> image = images.get(image_id)  # doctest: +SKIP
        >>> print(image.df)  # doctest: +SKIP
    """

    name = "images"

    def find(
        self,
        term: str | None = None,
        media_id: str | AND[str] | OR[str] | None = None,
        issue_id: str | AND[str] | OR[str] | None = None,
        is_front: bool | None = None,
        date_range: DateRange | None = None,
        visual_content: str | AND[str] | OR[str] | None = None,
        technique: str | AND[str] | OR[str] | None = None,
        communication_goal: str | AND[str] | OR[str] | None = None,
        content_type: str | AND[str] | OR[str] | None = None,
        embedding: Embedding | None = None,
        include_embeddings: bool = False,
        order_by: FindImagesOrderByLiteral | None = None,
        limit: int | None = None,
        offset: int | None = None,
    ) -> FindImagesContainer:
        """Find images in Impresso.

        Args:
            term: The search term for text-based search.
            media_id: Filter by newspaper media ID(s).
            issue_id: Filter by issue ID(s).
            is_front: Filter for front page images only.
            date_range: Filter by date range.
            visual_content: Filter by visual content category.
            technique: Filter by image technique.
            communication_goal: Filter by communication goal.
            content_type: Filter by content type.
            embedding: Image embedding for similarity search. Use `tools.embed_image()` or
                `tools.embed_text()` to generate embeddings from images.
            include_embeddings: Whether to include image embeddings in the response. Defaults to False.
            order_by: Sort order for results.
            limit: Maximum number of results to return per page. Defaults to 100.
            offset: Number of results to skip.

        Returns:
            FindImagesContainer: Data container with the first page of the search results.
        """
        page_limit = limit if limit is not None else DEFAULT_PAGE_SIZE
        embedding_with_limit = f"{embedding}:{page_limit}" if embedding else None

        filters: list[Filter] = []
        if media_id is not None:
            filters.extend(and_or_filter(media_id, "newspaper"))
        if issue_id is not None:
            filters.extend(and_or_filter(issue_id, "issue"))
        if is_front is not None:
            filters.append(Filter(type="isFront", daterange=None))
        if date_range is not None:
            filters.append(
                Filter(
                    type="daterange",
                    q=Q(DateRange._as_filter_value(date_range)),
                    context="exclude" if date_range.inverted else "include",
                    daterange=None,
                )
            )
        if visual_content is not None:
            filters.extend(and_or_filter(visual_content, "image_visual_content"))
        if technique is not None:
            filters.extend(and_or_filter(technique, "image_technique"))
        if communication_goal is not None:
            filters.extend(
                and_or_filter(communication_goal, "image_communication_goal")
            )
        if content_type is not None:
            filters.extend(and_or_filter(content_type, "image_content_type"))
        if embedding_with_limit is not None:
            filters.extend(and_or_filter(embedding_with_limit, "embedding"))

        filters_pb = filters_as_protobuf(filters or [])

        result = find_images.sync(
            client=self._api_client,
            term=term if term is not None else UNSET,
            order_by=(
                get_enum_from_literal(order_by, FindImagesOrderBy)
                if order_by is not None
                else UNSET
            ),
            limit=page_limit,
            offset=offset if offset is not None else UNSET,
            filters=filters_pb if filters_pb else UNSET,
            include_embeddings=include_embeddings if include_embeddings else UNSET,
        )
        raise_for_error(result)

        return FindImagesContainer(
            cast(FindImagesBaseFindResponse, result),
            FindImagesSchema,
            fetch_method=self.find,
            fetch_method_args={"term": term},
            web_app_search_result_url=_build_web_app_find_images_url(
                base_url=self._get_web_app_base_url(),
                term=term,
            ),
        )

    def get(
        self,
        id: str,
        include_embeddings: bool = False,
    ) -> GetImageContainer:
        """
        Get an image by its id.

        Args:
            id: The id of the image.
            include_embeddings: Whether to include embeddings in the response.

        Returns:
            GetImageContainer: The image data container.
        """
        result = get_image.sync(
            client=self._api_client,
            id=id,
            include_embeddings=include_embeddings,
        )
        raise_for_error(result)

        id_parts = id.split("-")
        issue_id = "-".join(id_parts[:-1])
        article_id = id_parts[-1]

        return GetImageContainer(
            result,
            Image,
            f"{self._get_web_app_base_url()}/issue/{issue_id}/view?articleId={article_id}",
        )

    def get_embeddings(self, id: str) -> list[str]:
        """
        Get the embeddings of an image by its id.

        Args:
            id: The id of the image.

        Returns:
            list[str]: The embeddings of the image if present (every embedding is returned
                in the canonical form: <model>:<base64_embedding>).
        """
        item = self.get(id, include_embeddings=True)
        return item.raw.get("embeddings", []) if item else []


def _build_web_app_find_images_url(
    base_url: str,
    term: str | None = None,
) -> str:
    query_params = {
        "q": term,
    }
    query_string = "&".join(
        f"{key}={value}" for key, value in query_params.items() if value is not None
    )
    url = f"{base_url}/images"
    return f"{url}?{query_string}" if query_string else url
