from __future__ import annotations

import asyncio
import logging
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._image = None

    def set_image(self, img: Optional[bytes]):
        self._image = img

    def get_image(self) -> Optional[bytes]:
        return self._image


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
        image = ex.get_image()
        return isinstance(image, bytes) and image.startswith(b"\xff\xd8") and image.find(b"JFIF\x00", 2, 24) != -1

    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}/picture", data=data, timeout=timeout
        ) as r:
            r.raise_for_status()
            ex.set_image(await r.read())

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