import asyncio
from io import BytesIO
from typing import TypeVar
from collections.abc import Awaitable

from PIL import Image, ImageDraw

from . import helper
from .helper import FontT, NodeType
from .source import BaseSource, EmojiCDNSource, HTTPBasedSource

T = TypeVar("T")
PILImage = Image.Image
PILDraw = ImageDraw.ImageDraw
ColorT = int | str | tuple[int, int, int] | tuple[int, int, int, int]


class Pilmoji:
    """The emoji rendering interface."""

    def __init__(
        self,
        *,
        source: BaseSource = EmojiCDNSource(),
        cache: bool = True,
        enable_tqdm: bool = True,
    ) -> None:
        self._cache: bool = cache
        self._source: BaseSource = source
        self._emoji_cache: dict[str, BytesIO] = {}
        self._discord_emoji_cache: dict[int, BytesIO] = {}

        if enable_tqdm:
            try:
                from tqdm.asyncio import tqdm

                self.__tqdm = tqdm
            except ImportError:
                self.__tqdm = None

    async def aclose(self) -> None:
        if isinstance(self._source, HTTPBasedSource):
            await self._source.aclose()

    async def _get_emoji(self, emoji: str) -> BytesIO | None:
        if self._cache and emoji in self._emoji_cache:
            return self._emoji_cache[emoji]

        if bytesio := await self._source.get_emoji(emoji):
            if self._cache:
                self._emoji_cache[emoji] = bytesio
            return bytesio

        return None

    async def _get_discord_emoji(self, id: int) -> BytesIO | None:
        if self._cache and id in self._discord_emoji_cache:
            return self._discord_emoji_cache[id]

        if bytesio := await self._source.get_discord_emoji(id):
            if self._cache:
                self._discord_emoji_cache[id] = bytesio
            return bytesio

        return None

    def _resize_emoji(self, bytesio: BytesIO, size: float) -> PILImage:
        """Resize emoji to fit the font size"""
        bytesio.seek(0)
        with Image.open(bytesio).convert("RGBA") as emoji_img:
            emoji_size = int(size) - 1
            aspect_ratio = emoji_img.height / emoji_img.width
            return emoji_img.resize(
                (emoji_size, int(emoji_size * aspect_ratio)),
                Image.Resampling.LANCZOS,
            )

    async def text(
        self,
        image: PILImage,
        xy: tuple[int, int],
        text: str,
        font: FontT,
        fill: ColorT | None = None,
    ) -> None:
        """Simplified text rendering method with Unicode emoji support.

        Parameters
        ----------
        image: PILImage
            The image to render onto
        xy: tuple[int, int]
            Rendering position (x, y)
        text: str
            The text to render (supports single or multiple lines)
        font: FontT
            The font to use
        fill: ColorT | None
            Text color, defaults to black
        """
        if not text:
            return

        x, y = xy
        draw = ImageDraw.Draw(image)
        line_height = helper.get_font_height(font)

        # check text has emoji
        if not helper.has_emoji(text):
            for line in text.splitlines():
                draw.text((x, y), line, font=font, fill=fill)
                y += line_height
            return

        # Parse text into nodes (Unicode emoji only)
        lines = helper.to_nodes(text)

        # Collect all unique Unicode emojis to download
        emj_set = {
            node.content
            for line in lines
            for node in line
            if node.type is NodeType.EMOJI
        }

        # Download all emojis concurrently
        emjios = await self.gather(
            *[self._get_emoji(emoji) for emoji in emj_set],
        )
        emj_map = dict(zip(emj_set, emjios))

        # Render each line
        font_size = helper.get_font_size(font)
        y_diff = int((line_height - font_size) / 2)

        # Pre-resize emojis
        resized_emojis: dict[str, PILImage] = {}
        for emoji, bytesio in emj_map.items():
            if bytesio:
                resized_emojis[emoji] = self._resize_emoji(bytesio, font_size)

        for line in lines:
            cur_x = x

            for node in line:
                if node.type is NodeType.EMOJI:
                    if emoji_img := resized_emojis.get(node.content):
                        image.paste(emoji_img, (cur_x, y + y_diff), emoji_img)
                    else:
                        draw.text((cur_x, y), node.content, font=font, fill=fill)
                    cur_x += int(font_size)
                else:
                    # Text node
                    draw.text((cur_x, y), node.content, font=font, fill=fill)
                    cur_x += int(font.getlength(node.content))

            y += line_height

    async def text_with_ds_emj(
        self,
        image: PILImage,
        xy: tuple[int, int],
        text: str,
        font: FontT,
        fill: ColorT | None = None,
    ) -> None:
        """Simplified text rendering method with Unicode and Discord emoji support.

        Parameters
        ----------
        image: PILImage
            The image to render onto
        xy: tuple[int, int]
            Rendering position (x, y)
        text: str
            The text to render (supports single or multiple lines)
        font: FontT
            The font to use
        fill: ColorT | None
            Text color, defaults to black
        """
        if not text:
            return

        x, y = xy
        draw = ImageDraw.Draw(image)
        line_height = helper.get_font_height(font)

        if not helper.has_emoji(text, False):
            for line in text.splitlines():
                draw.text((x, y), line, font=font, fill=fill)
                y += line_height
            return

        # Parse text into nodes
        lines = helper.to_nodes(text, False)

        # Collect all unique emojis to download
        emj_set = {
            node.content
            for line in lines
            for node in line
            if node.type is NodeType.EMOJI
        }
        ds_emj_set = {
            int(node.content)
            for line in lines
            for node in line
            if node.type is NodeType.DISCORD_EMOJI
        }

        # Download all emojis concurrently
        emjios = await self.gather(
            *[self._get_emoji(emoji) for emoji in emj_set],
            *[self._get_discord_emoji(eid) for eid in ds_emj_set],
        )

        emj_num = len(emj_set)
        emoji_results = emjios[:emj_num]
        discord_results = emjios[emj_num:]

        # Build emoji mappings
        emj_map = dict(zip(emj_set, emoji_results))
        ds_emj_map = dict(zip(ds_emj_set, discord_results))

        # Render each line
        font_size = helper.get_font_size(font)
        y_diff = int((line_height - font_size) / 2)

        # Pre-resize emojis
        resized_emojis: dict[str | int, PILImage] = {}
        for emoji, bytesio in emj_map.items():
            if bytesio:
                resized_emojis[emoji] = self._resize_emoji(bytesio, font_size)
        for eid, bytesio in ds_emj_map.items():
            if bytesio:
                resized_emojis[eid] = self._resize_emoji(bytesio, font_size)

        for line in lines:
            cur_x = x

            for node in line:
                emoji_img = None

                match node.type:
                    case NodeType.EMOJI:
                        emoji_img = resized_emojis.get(node.content)
                    case NodeType.DISCORD_EMOJI:
                        emoji_img = resized_emojis.get(int(node.content))

                # Render emoji or text
                if emoji_img:
                    image.paste(emoji_img, (cur_x, y + y_diff), emoji_img)
                    cur_x += int(font_size)
                else:
                    draw.text((cur_x, y), node.content, font=font, fill=fill)
                    cur_x += int(font.getlength(node.content))

            y += line_height

    async def gather(self, *tasks: Awaitable[T]) -> list[T]:
        if self.__tqdm is None:
            return await asyncio.gather(*tasks)

        return await self.__tqdm.gather(
            *tasks,
            desc="Fetching Emojis",
            colour="green",
            bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}",
        )

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_value, traceback) -> None:
        await self.aclose()

    def __repr__(self) -> str:
        return f"<Pilmoji source={self._source} cache={self._cache}>"
