from __future__ import annotations

import asyncio
import logging
import struct
from typing import Optional

import aiohttp

from clarabot_nano.indexer.extractor.base import (
    AbstractExtraction,
    AbstractContentExtractor,
    ExtractError,
    AbstractExtractionSource,
)

logger = logging.getLogger(__name__)


class PicotureExtraction(AbstractExtraction):
    def __init__(self):
        self._preview_size = None
        self._preview = None
        self._thumbnail = None

    def set_result(self, data: Optional[bytes]):
        if not data or len(data) < 4:
            self._preview = self._thumbnail = None
        else:
            offset = struct.unpack(">i", data[:4])[0] + 4
            t = data[4:offset]
            p_size = struct.unpack(">i", data[offset : offset + 4])[0]
            p = data[offset + 4 :]
            if self.is_image_valid(t) and self.is_image_valid(p) and 0 < p_size <= 1073741823:
                self._thumbnail = t
                self._preview = p
                self._preview_size = p_size
            else:
                self._preview_size = self._preview = self._thumbnail = None

    def get_preview_size(self) -> Optional[int]:
        return self._preview_size

    def get_preview(self) -> Optional[bytes]:
        return self._preview

    def get_thumbnail(self) -> Optional[bytes]:
        return self._thumbnail

    @classmethod
    def is_image_valid(cls, image: Optional[bytes]) -> bool:
        return isinstance(image, bytes) and image.startswith(b"\xff\xd8") and image.find(b"JFIF\x00", 2, 24) != -1


class PicotureContentExtractor(AbstractContentExtractor):
    """
    Extract a small image (thumbnail) from a file as well as possible.
    """

    CHECK_PNG = (
        b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde"
        b"\x00\x00\x00\x0cIDAT\x08\xd7cxy\x9c\x0f\x00\x04[\x01\xbf\xdc\x10^\x0b\x00\x00\x00\x00IEND\xaeB`\x82"
    )

    def _create_extraction(self, source: AbstractExtractionSource) -> AbstractExtraction:
        return PicotureExtraction()

    async def check_operable(self, timeout: float = 5.0) -> bool:
        ex = PicotureExtraction()
        try:
            await self._extract_content(self.CHECK_PNG, ex, timeout=timeout)
        except asyncio.TimeoutError:
            logger.warning("Extractor check timed out")
            return False
        except (aiohttp.ClientError, ExtractError) as e:
            logger.warning("Extractor check failed: %s", e)
            return False
        return ex.get_preview() is not None  # checking one is sufficient, either none or both are set

    async def _extract_content(self, data, ex: PicotureExtraction, timeout: float = 60.0):
        async with self.session.post(
            f"http://127.0.0.1:{self.ec.picoture_port}/pictures", data=data, timeout=timeout
        ) as r:
            r.raise_for_status()
            ex.set_result(await r.read())

    def _is_status_successful(self, source: AbstractExtractionSource, ex: AbstractExtraction, status: int) -> bool:
        return status == 422 or status == 415  # unsupported file format
