from __future__ import annotations

import asyncio
import logging
import os
import re
import sys
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Optional, Union

import aiofiles
import aiohttp
from aiofiles.threadpool.binary import AsyncBufferedReader

from clarabot_nano.indexer.config import IndexerConfig
from clarabot_nano.indexer.extractor.base import (
    AbstractExtractionSource,
    AbstractContentExtractor,
    AbstractExtraction,
    ExtractError,
)
from clarabot_nano.indexer.extractor.send import SendEmptyFileError, SendFileHandler, send_open_file, send_memory_view
from clarabot_nano.utility.mime import nano_mimetypes

logger = logging.getLogger(__name__)


class MimeTypeExtraction(AbstractExtraction):
    def __init__(self):
        self.mimetype = None


class MimeTypeExtractionSource(AbstractExtractionSource):
    def __init__(self, name: str, path: Path, head_chunk: bytes):
        self.name = name
        self.path = path
        self.head_chunk = head_chunk

    async def _data(self):
        yield self.head_chunk

    def fallback_summary(self) -> str:
        return self.name

    @asynccontextmanager
    async def open(self):
        yield SendFileHandler(self.name, None, self._data())

    def describe(self) -> str:
        return str(self.path)


class FileExtractionSource(AbstractExtractionSource, AbstractContentExtractor):
    """
    This object prepares information about a file source for extraction and handles data buffering.
    It implements a Tika extractor to detect the mimetype of the file by magic-byte recognition.
    Depending on the information the content of the file may need multiple work passes. Because of this, as long as
    there is enough memory to hold the file it may be much more optimal to hold it for all the passes instead of
    always reading it from the storage. The IO latency of the storage may be high if we're looking at some NAS or
    share anyways. Rarely, even extraction attempts may fail and would require re-reading the contents for the same
    type of processing, which can also be improved if we use memory buffering. For now we only pre-load contents
    into memory, when we know multiple passes will take place.
    """

    # TODO: Some metadata that is extracted from some of these, but we need custom handling and they are low priority.
    SKIP_EXTRACT_TEXT_RE = re.compile(r"^(image|audio|video)/.+$", re.IGNORECASE)
    # Optimized for what Picoture supports.
    EXTRACT_IMAGE_RE = re.compile(
        r"^(application/pdf|image/(?:"
        r"jpeg|png|gif|(?:x-(?:ms-)?)?bmp|vnd\.wap\.wbmp|"  # ImageIO support in Java 8
        r"svg\+xml|"
        r"(?:x-)?wmf|vnd\.microsoft\.icon|"  # Windows metafile and icon
        r"vnd\.radiance|"  # HDR image
        r"icns|"  # Apple icon
        r"vnd\.zbrush\.[pd]cx|"  # ZSoft PCX & DCX
        r"x-pict|"  # Apple Picture
        r"x-portable-(?:any|bit|gray|pix)map|"  # PNM
        r"vnd\.adobe\.photoshop|x-psd|"  # Photoshop
        r"x-rgb|"  # SGI
        r"x-t(?:ar)?ga|"
        r"tiff|"
        r"webp"
        r"))$",
        re.IGNORECASE,
    )
    JS_EXT_RE = re.compile(r".+\.js$", re.IGNORECASE)
    HDR_EXT_RE = re.compile(r".+\.(?:pic|hdr|rgbe|xyze|rad|img)$", re.IGNORECASE)
    HDR_SIGNATURE_RE = re.compile(rb"^#\?(?:RGBE|RADIANCE)\n")

    def __init__(
        self,
        session: aiohttp.ClientSession,
        ec: IndexerConfig,
        name: str,
        path: Path,
        s: os.stat_result,
        until: int = sys.maxsize,
        chunk_size: int = 65536,
    ):
        AbstractContentExtractor.__init__(self, session, ec)
        self.name = name
        self.path = path
        self.stat = s
        self.until = until  # read at least until this point (or the end of the file)
        self.chunk_size = chunk_size
        self.content: Optional[Union[bytearray, AsyncBufferedReader]] = None
        self.content_type: Optional[str] = None
        self.extract_text = True
        self.extract_image = True

    def _set_content_type(self, mimetype: Optional[str], head_chunk: bytes):
        """
        Evaluate the detected mimetype by Tika.
        """
        if len(head_chunk) < 128:
            # Magic detection often fails for tiny files, so we naively try to improve it.
            if mimetype == "application/octet-stream":
                if len(head_chunk.decode("utf-8", errors="ignore")) >= len(head_chunk) * 0.99:
                    mimetype = "text/plain"
        elif self.HDR_SIGNATURE_RE.match(head_chunk) is not None and self.HDR_EXT_RE.match(self.name) is not None:
            # The Radiance HDR format is not recognized by Tika, so we do it ourselves.
            mimetype = "image/vnd.radiance"
        elif mimetype == "application/javascript":
            # There is no magic to detect Javascript and it often mistakes related files, so we try to improve it.
            if self.JS_EXT_RE.match(self.name) is None:
                mimetype = nano_mimetypes.guess_type(self.name, strict=False)[0]
        self.content_type = mimetype
        self.extract_text = bool(
            head_chunk and (not self.content_type or self.SKIP_EXTRACT_TEXT_RE.match(self.content_type) is None)
        )
        self.extract_image = bool(
            head_chunk and (not self.content_type or self.EXTRACT_IMAGE_RE.match(self.content_type) is not None)
        )

    @asynccontextmanager
    async def prepare(self, eager_buffer_limit: int = 2097152):
        """
        For optimal file content processing and IO usage, this is a little complex.
        We will stream the file data unless multiple passes are required for extraction, while reading into memory
        the first chunk (head) of the file for mimetype detection with magic-byte recognition.
        """
        try:
            async with aiofiles.open(self.path, "rb") as f:
                head_chunk: bytes = await f.read(self.chunk_size)
                mimetype_extraction: MimeTypeExtraction = await self.extract(
                    MimeTypeExtractionSource(self.name, self.path, head_chunk)
                )
                self._set_content_type(mimetype_extraction.mimetype, head_chunk)
                buffer = self.stat.st_size < eager_buffer_limit or (self.extract_text and self.extract_image)
                if buffer:
                    data = bytearray()
                    async for chunk in send_open_file(f, head_chunk, self.until, self.chunk_size):
                        data.extend(chunk)
                    self.content = data
                    # fallthrough to the yield outside of the file handle
                else:
                    self.content = f
                    yield self
            if buffer:
                yield self
        finally:
            self.content = None

    def fallback_summary(self) -> str:
        return self.name

    @asynccontextmanager
    async def open(self):
        if self.content is None:  # pragma: no cover
            raise Exception("Must be prepared when submitting for extraction")
        elif isinstance(self.content, AsyncBufferedReader):
            # read the file without buffering
            await self.content.seek(0)
            chunk = await self.content.read(self.chunk_size)
            if not chunk:
                raise SendEmptyFileError()
            yield SendFileHandler(
                self.path.name, self.content_type, send_open_file(self.content, chunk, self.until, self.chunk_size)
            )
        else:
            assert isinstance(self.content, bytearray)
            # When the data-buffer is bytes, it is the first chunk (head of the file)
            if not len(self.content):
                raise SendEmptyFileError()
            v = memoryview(self.content)
            try:
                yield SendFileHandler(
                    self.path.name, self.content_type, send_memory_view(v, self.until, self.chunk_size)
                )
            finally:
                v.release()  # explicit free

    def describe(self) -> str:
        return str(self.path)

    # extractor api for mimetype detection
    def _create_extraction(self, source: AbstractExtractionSource) -> MimeTypeExtraction:
        return MimeTypeExtraction()

    async def check_operable(self, timeout: float = 5.0) -> bool:
        try:
            await self._extract_content(b"", MimeTypeExtraction(), 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 True

    async def _extract_content(self, data, ex: MimeTypeExtraction, timeout: float = 5.0):
        async with self.session.put(
            f"http://127.0.0.1:{self.ec.tika_port}/detect/stream", data=data, timeout=timeout
        ) as r:
            r.raise_for_status()
            ex.mimetype = (await r.text()).strip()
