import asyncio as aio
import mimetypes
import random
import sys
import time
from abc import ABC, abstractmethod
from collections.abc import AsyncIterable, Callable
from math import floor
from pathlib import Path
from typing import Generic, NamedTuple, ParamSpec, TypeAlias, TypedDict, TypeVar
from typing_extensions import override

from cookit.common import race
from cookit.loguru import warning_suppress
from httpx import AsyncClient, Response
from nonebot import get_driver, logger

from .config import BG_PRELOAD_CACHE_DIR, DEFAULT_BG_PATH, config

if sys.version_info >= (3, 11):
    from asyncio.taskgroups import TaskGroup
else:
    from taskgroup import TaskGroup


class BgBytesData(NamedTuple):
    data: bytes | None
    mime: str


class BgFileData(NamedTuple):
    path: Path | None
    mime: str


BgData: TypeAlias = BgBytesData | BgFileData

BGProviderType = Callable[[int], AsyncIterable[BgData]]
T = TypeVar("T")
TBP = TypeVar("TBP", bound=BGProviderType)
P = ParamSpec("P")

DEFAULT_MIME = "application/octet-stream"

registered_bg_providers: dict[str, BGProviderType] = {}


def get_bg_files() -> list["Path"]:
    if not config.ps_bg_local_path.exists():
        logger.warning("Custom background path does not exist, fallback to default")
        return [DEFAULT_BG_PATH]
    if config.ps_bg_local_path.is_file():
        return [config.ps_bg_local_path]

    files = [x for x in config.ps_bg_local_path.glob("*") if x.is_file()]
    if not files:
        logger.warning("Custom background dir has no file in it, fallback to default")
        return [DEFAULT_BG_PATH]
    return files


BG_FILES = get_bg_files()


def refresh_bg_files():
    global BG_FILES
    BG_FILES = get_bg_files()


def bg_provider(name: str | None = None):
    def deco(func: TBP) -> TBP:
        provider_name = name or func.__name__
        if provider_name in registered_bg_providers:
            raise ValueError(f"Duplicate bg provider name `{provider_name}`")
        registered_bg_providers[provider_name] = func
        return func

    return deco


def iter_batch_sizes(size: int, max_size: int):
    if size <= max_size:
        yield size
    else:
        full_sizes = floor(max_size / size)
        for _ in range(full_sizes):
            yield max_size
        if rest_count := full_sizes * max_size:
            yield rest_count


def resp_to_bg_data(resp: Response):
    return BgBytesData(
        resp.content,
        (resp.headers.get("Content-Type") or DEFAULT_MIME),
    )


class CoIterator(ABC, Generic[T]):
    def __init__(self):
        self.queue = aio.Queue[T | None]()

    @abstractmethod
    async def run_tasks(self): ...

    async def run(self):
        await self.run_tasks()
        await self.queue.put(None)

    async def __aiter__(self):
        async with TaskGroup() as t:
            t.create_task(self.run())
            while (x := await self.queue.get()) is not None:
                yield x


@bg_provider("loli")
class LoliBGProvider(CoIterator[BgData]):
    def __init__(self, num: int):
        super().__init__()
        self.num = num
        self.sem = aio.Semaphore(4)

    async def task_piece(self, cli: AsyncClient):
        async with self.sem:
            with warning_suppress("Failed to fetch image"):
                x = resp_to_bg_data(
                    (
                        await cli.get("https://www.loliapi.com/acg/pe/")
                    ).raise_for_status(),
                )
                await self.queue.put(x)

    @override
    async def run_tasks(self):
        async with AsyncClient(
            follow_redirects=True,
            proxy=config.proxy,
            timeout=config.ps_req_timeout,
        ) as cli:
            await aio.gather(*(self.task_piece(cli) for _ in range(self.num)))


class LoliconRespDataUrls(TypedDict):
    original: str


class LoliconRespData(TypedDict):
    urls: LoliconRespDataUrls


class LoliconResp(TypedDict):
    data: list[LoliconRespData]


@bg_provider("lolicon")
class LoliconBGProvider(CoIterator[BgData]):
    def __init__(self, num: int):
        super().__init__()
        self.num = num
        self.sem = aio.Semaphore(4)
        self.url_queue = aio.Queue[str | None]()

    async def do_fetch_urls_piece(self, num: int, cli: AsyncClient):
        with warning_suppress("Failed to fetch urls"):
            resp = await cli.get(
                "https://api.lolicon.app/setu/v2",
                params={
                    "num": num,
                    "r18": config.ps_bg_lolicon_r18_type,
                    "proxy": "false",
                    "excludeAI": "true",
                },
            )
            data: LoliconResp = resp.raise_for_status().json()
            for x in data["data"]:
                await self.url_queue.put(x["urls"]["original"])

    async def fetch_urls_task_f(self):
        async with AsyncClient(
            follow_redirects=True,
            proxy=config.proxy,
            timeout=config.ps_req_timeout,
        ) as cli:
            for x in iter_batch_sizes(self.num, 20):
                await self.do_fetch_urls_piece(x, cli)
        await self.url_queue.put(None)

    async def fetch_image(self, url: str, cli: AsyncClient):
        async with self.sem:
            with warning_suppress("Failed to fetch image"):
                bg = resp_to_bg_data((await cli.get(url)).raise_for_status())
                await self.queue.put(bg)

    @override
    async def run_tasks(self):
        pixiv_client = AsyncClient(
            follow_redirects=True,
            proxy=config.proxy,
            timeout=config.ps_req_timeout,
            headers={
                "User-Agent": (
                    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                    "AppleWebKit/537.36 (KHTML, like Gecko) "
                    "Chrome/119.0.0.0 "
                    "Safari/537.36"
                ),
                "Referer": "https://www.pixiv.net/",
            },
        )
        async with TaskGroup() as t, pixiv_client:
            t.create_task(self.fetch_urls_task_f())
            while (x := await self.url_queue.get()) is not None:
                t.create_task(self.fetch_image(x, pixiv_client))


@bg_provider()
async def local(num: int):
    files = random.sample(BG_FILES, num)
    # logger.debug(f"Chosen background `{files}`")
    for x in files:
        yield BgFileData(
            x,
            mimetypes.guess_type(x)[0] or DEFAULT_MIME,
        )


def create_none_bg():
    return BgBytesData(None, DEFAULT_MIME)


@bg_provider()
async def none(num: int):
    for _ in range(num):
        yield create_none_bg()


async def fetch_bg(num: int) -> AsyncIterable[BgData]:
    if config.ps_bg_provider not in registered_bg_providers:
        logger.warning(
            f"Unknown background provider `{config.ps_bg_provider}`, fallback to local",
        )
        async for x in local(num):
            yield x
        return

    # at least we should return one image (x)
    # has_img = False
    try:
        provider = registered_bg_providers[config.ps_bg_provider]
        async for x in provider(num):
            # has_img = True
            yield x
    except Exception:
        logger.exception(
            "Error when getting background, fallback to get one local bg",
        )
        async for x in local(1):
            yield x
    # else:
    #     if has_img:
    #         return
    #     logger.warning(
    #         "Background provider returned empty iterator, fallback to get one local bg",
    #     )
    #     async for x in local(1):
    #         yield x


def cache_bg(bg: BgBytesData):
    if not bg.data:
        return BgFileData(None, bg.mime)
    BG_PRELOAD_CACHE_DIR.mkdir(parents=True, exist_ok=True)
    path = BG_PRELOAD_CACHE_DIR / f"{time.time_ns()}.{bg.mime.split('/')[-1]}"
    path.write_bytes(bg.data)
    return BgFileData(path, bg.mime)


def read_cached_bg_file(bg: BgFileData) -> BgBytesData | None:
    if not bg.path:
        return BgBytesData(None, bg.mime)
    with warning_suppress("Failed to read cached file"):
        data = bg.path.read_bytes()
        if bg.path.is_relative_to(BG_PRELOAD_CACHE_DIR):
            with warning_suppress("Failed to unlink cached file"):
                bg.path.unlink()
        return BgBytesData(data, bg.mime)
    return None


async def get_one_fallback() -> BgBytesData:
    with warning_suppress("Failed to get local bg file, fallback to none"):
        async for x in local(1):
            if bg := read_cached_bg_file(x):
                return bg
    logger.warning("Failed to read local bg file, fallback to none")
    return create_none_bg()


class BgPreloader:
    def __init__(self, preload_count: int):
        # if preload_count < 1:
        #     raise ValueError("preload_count must be greater than or equals 1")
        self.preload_count = preload_count
        self.background_queue = aio.Queue[BgData]()
        self.current_load_task_main: aio.Task | None = None
        self.consumed_in_loading: bool = False
        self.image_got_signal = aio.Event()
        self.fire_tasks: set[aio.Task] = set()

    # we allow fetch_bg return less image than we require
    async def preload_task(
        self,
        count: int,
        fire: bool = False,
        fire_done_signal: aio.Event | None = None,
    ):
        logger.debug(f"Preload task started, will preload {count} images, {fire=}")
        try:
            async for x in fetch_bg(count):
                logger.debug("Got one image")
                if self.preload_count > 0 or (
                    fire_done_signal and fire_done_signal.is_set()
                ):
                    x = cache_bg(x) if isinstance(x, BgBytesData) else x
                await self.background_queue.put(x)
                self.image_got_signal.set()
                self.image_got_signal.clear()
        except Exception:
            logger.exception("Unexpected error occurred in preload task")
        else:
            logger.debug("Preload task finished")

        if fire:
            return
        if (
            self.consumed_in_loading
            or self.background_queue.qsize() < self.preload_count
        ):
            self.consumed_in_loading = False
            self.start_preload()
        else:
            self.current_load_task_main = None

    def start_preload(self, force: bool = False):
        count = self.preload_count - self.background_queue.qsize()
        if count <= 0 and not force:
            logger.debug(
                "Current background queue size meets preload count, skip preload",
            )
            return
        task = aio.create_task(self.preload_task(count))
        self.current_load_task_main = task

    def set_defer_preload(self):
        if self.current_load_task_main:
            logger.debug("Main preload task already running, set flag")
            self.consumed_in_loading = True
        else:
            self.start_preload()

    async def _get_on_fire(self) -> BgBytesData:
        task_done_signal = aio.Event()
        fire_task = aio.create_task(
            self.preload_task(1, fire=True, fire_done_signal=task_done_signal),
        )
        fire_task.add_done_callback(lambda _: task_done_signal.set())
        fire_task.add_done_callback(lambda _: self.fire_tasks.discard(fire_task))
        self.fire_tasks.add(fire_task)
        try:
            await race(
                # self.image_got_signal.wait(),  # lazy to handle this racing condition now
                task_done_signal.wait(),
                aio.sleep(15),
            )
        finally:
            task_done_signal.set()
            # fire_task.cancel()  # should we cancel here? i'm letting it cache to queue

        if not self.background_queue.empty():
            bg = await self.background_queue.get()
            self.set_defer_preload()
            if (not isinstance(bg, BgFileData)) or (bg := read_cached_bg_file(bg)):
                return bg

        logger.error("Unable to get an background image, falling back to local")
        return await get_one_fallback()

    async def get(self) -> BgBytesData:
        self.set_defer_preload()

        while not self.background_queue.empty():
            bg = await self.background_queue.get()
            self.set_defer_preload()
            if (not isinstance(bg, BgFileData)) or (bg := read_cached_bg_file(bg)):
                return bg

        # normally all items in queue should be valid
        # if they not, we should fetch
        return await self._get_on_fire()


bg_preloader = BgPreloader(config.ps_bg_preload_count)

driver = get_driver()


@driver.on_shutdown
async def _():
    for t in bg_preloader.fire_tasks:
        t.cancel()