from __future__ import annotations

import asyncio
import errno
import itertools
import logging
import mimetypes
import sys
import time
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Optional, List, Any, AsyncIterable

import aiofiles
import aiohttp
from aiofiles.threadpool.binary import AsyncBufferedIOBase
from lxml import etree

from clarabot_nano.indexer.extractor.base import (
    AbstractCollector,
    AbstractExtraction,
    CollectorHandler,
    SPACE_SEARCH_RE,
    NEWLINE_STRIP_RE,
    AbstractContentExtractor,
    IgnoredMimeType,
    ExtractError,
    EnoughText,
    ExtractorInoperable,
    RetryExtractError,
    ServerInternalExtractError,
)
from clarabot_nano.indexer.extractor.body import BodyHandler
from clarabot_nano.indexer.extractor.head import TitleHandler, HeadHandler
from clarabot_nano.utility.path.error import NtOSError

logger = logging.getLogger(__name__)


class Extraction(AbstractCollector, AbstractExtraction):
    """
    Process XML made by Tika to create documents for indexing.
    """

    def __init__(self, ignored_mimes: set, extract_text_limit: int, fallback_summary: str):
        self.fallback_summary = fallback_summary
        # Collector handlers. As simple as possible with extensibility.
        self._title_handler = TitleHandler()
        self._head_handler = HeadHandler(ignored_mimes, {"title": self._title_handler})
        self._body_handler = BodyHandler(extract_text_limit)
        self._handler_stack: List[CollectorHandler] = [
            CollectorHandler({"head": self._head_handler, "body": self._body_handler}),
        ]

    def start(self, tag: str, attrs: dict):
        tag = tag.rpartition("}")[-1]  # we don't care about namespaces
        handler = self._handler_stack[-1]
        if handler.sub_handlers and (h := handler.sub_handlers.pop(tag, None)):
            self._handler_stack.append(h)
            h.init(attrs)
        else:
            handler.start(tag, attrs)

    def end(self, tag: str):
        tag = tag.rpartition("}")[-1]  # we don't care about namespaces
        handler = self._handler_stack[-1]
        if handler.TAG == tag:
            self._handler_stack.pop().close()
        else:
            handler.end(tag)

    def data(self, text: str):
        self._handler_stack[-1].data(text)

    def comment(self, text: str):
        self._handler_stack[-1].comment(text)

    def close(self):
        while self._handler_stack:
            self._handler_stack.pop().close()

    def get_summary(self) -> str:
        summary = []
        if t := self._title_handler.title:
            summary.append(t)
        if hs := self._body_handler.headers:
            if hs[-1] == " ":
                hs = itertools.islice(hs, len(hs) - 1)
            summary.append("".join(hs))
        summary.append(self.fallback_summary)
        text = "\n".join(summary)
        text = SPACE_SEARCH_RE.sub(" ", NEWLINE_STRIP_RE.sub("\n", text)).strip()
        return text

    def get_content(self) -> Optional[str]:
        texts = self._body_handler.texts
        if texts and texts[-1] == " ":
            texts = itertools.islice(texts, len(texts) - 1)
        text = "".join(texts)
        return SPACE_SEARCH_RE.sub(" ", NEWLINE_STRIP_RE.sub("\n", text)).strip()


class CollectorXMLParser(etree.XMLParser):
    def __init__(self, collector: AbstractCollector):
        etree.XMLParser.__init__(
            self,
            encoding="utf-8",
            attribute_defaults=False,
            dtd_validation=False,
            load_dtd=False,
            no_network=True,
            ns_clean=False,
            recover=True,
            schema=None,
            huge_tree=False,
            remove_blank_text=True,
            resolve_entities=True,
            remove_comments=True,
            remove_pis=True,
            strip_cdata=True,
            collect_ids=False,
            target=collector,
            compact=True,
        )

    def feed_end(self):
        pass


class SendEmptyFileError(Exception):
    pass


async def send_open_file(f: AsyncBufferedIOBase, chunk: bytes, until: int, chunk_size: int):
    """
    Buffered file sender for extraction, with maximum size control.
    """
    yield chunk
    pos = len(chunk)
    chunk = await f.read(chunk_size)
    while chunk:
        yield chunk
        pos += len(chunk)
        if pos >= until:
            break
        chunk = await f.read(chunk_size)


class SendFileHandler(object):
    __slots__ = ("path", "iter")

    @classmethod
    @asynccontextmanager
    async def create(cls, path: Path, until: int = sys.maxsize, chunk_size: int = 65536):
        """
        Open the file and check if it's empty before here. This avoids connecting to the extractor for inaccessible
        (PermissionError for example) or empty files.
        """
        async with aiofiles.open(path, "rb") as f:
            chunk = await f.read(chunk_size)
            if not chunk:
                raise SendEmptyFileError()
            yield cls(path, send_open_file(f, chunk, until, chunk_size))

    def __init__(self, path: Path, iter_: AsyncIterable[bytes]):
        self.path = path
        self.iter = iter_


class SendFilePayload(aiohttp.AsyncIterablePayload):
    def __init__(self, value: SendFileHandler, *args: Any, **kwargs: Any) -> None:
        if not isinstance(value, SendFileHandler):
            raise TypeError(f"value argument must be SendFileHandle, got {type(value)!r}")

        kwargs["filename"] = value.path.name
        kwargs["content_type"] = None  # use mimetype detection from filename

        aiohttp.AsyncIterablePayload.__init__(self, value.iter, *args, **kwargs)


aiohttp.PAYLOAD_REGISTRY.register(SendFilePayload, SendFileHandler)

# force loading the mime db, before applying fixes
mimetypes.guess_type("test.js")
# Do not allow overriding the standard mime-types at the least. Such should not happen, and we've observed just
# opening a .js file in notepad on windows modifying the registry associations and stating that .js files are
# text/plain, which is not useful and usually incorrect.
for ext_, type_ in mimetypes._types_map_default.items():
    mimetypes._db.add_type(type_, ext_)


class FileContentExtractor(AbstractContentExtractor):
    """
    Extract plain-text contents from a file as well as possible.
    """

    CHECK_TEXT = "A little text to test the extractor."
    TIKA_HEADERS = {
        "Accept": "text/xml",
        # https://tika.apache.org/1.26/api/org/apache/tika/parser/pdf/PDFParserConfig.html
        "X-Tika-PDFSortByPosition": "true",
        "X-Tika-PDFEnableAutoSpace": "true",
        # What would this do exactly?
        # "X-Tika-PDFMaxMainMemoryBytes": "300000000",
    }

    async def extract(self, path: Path) -> AbstractExtraction:
        ex = self._create_extraction(self.ec.indexer_ignored_mimes, self.ec.indexer_extract_text_limit, path.name)
        try:
            # NOTE: Tika does not handle unfinished files gracefully (may depend on file-type).
            async with SendFileHandler.create(path) as payload:
                await self._extract_xml_content(payload, ex)
        except SendEmptyFileError:
            pass  # empty file is not sent because there's nothing to extract
        except (IgnoredMimeType, EnoughText):
            pass  # intentional interrupts, return the extraction
        except aiohttp.ClientResponseError as e:
            try:
                await self.ensure_operable()
            except Exception as e2:
                raise e2 from e  # chain the exceptions for clarity
            if e.status >= 500:
                # This could be a bad
                raise ServerInternalExtractError(ex) from e
            if e.status != 422:  # 422 is returned for empty files; so it's a success...
                # Other error codes in the 4xx range are likely related to the HTTP API server, they may be temporary?
                raise RetryExtractError(ex) from e
        except aiohttp.ClientError as e:  # Something infrastructural: payload parsing or connection error most likely.
            try:
                await self.ensure_operable()
            except Exception as e2:
                raise e2 from e  # chain the exceptions for clarity
            # These issues should be temporary.
            raise RetryExtractError(ex) from e
        except asyncio.TimeoutError as e:
            try:
                await self.ensure_operable()
            except Exception as e2:
                raise e2 from e  # chain the exceptions for clarity
            # The extractor should not get stuck indefinitely, so the machine may be overwhelmed.
            raise RetryExtractError(ex)
        except etree.LxmlError:
            logger.exception("Failed to parse Tika response for: %s", path)
        except NtOSError as e:
            # NOTE: We treat some OS errors as final even though their reason may be varied. For example:
            #   - Google-docs links in a mounted Google Drive raise [Errno 22]/EINVAL for reading.
            if e.errno != errno.EINVAL:
                raise e
        return ex

    async def ensure_operable(self, timeout: float = 30.0):
        try:
            t = time.perf_counter() + timeout
            while True:
                attempt_start = time.perf_counter()
                attempt_timeout = max(0.5, min(t - attempt_start, 5.0))
                if await self.check_operable(timeout=attempt_timeout):
                    break
                if time.perf_counter() > t:
                    raise ExtractorInoperable()
                # simple backoff from retrying
                await asyncio.sleep(max(0.5, attempt_timeout - (time.perf_counter() - attempt_start)))
        except ExtractError as e:
            raise ExtractorInoperable() from e

    async def check_operable(self, timeout: float = 5.0) -> bool:
        ex = self._create_extraction(set(), 65536, "fallback")
        try:
            await self._extract_xml_content(self.CHECK_TEXT.encode("utf-8"), ex, timeout=timeout)
        except asyncio.TimeoutError:
            logger.warning("Extractor check timed out")
            return False
        except (aiohttp.ClientError, etree.LxmlError, ExtractError) as e:
            logger.warning("Extractor check failed: %s", e)
            return False
        return ex.get_content().find(self.CHECK_TEXT) >= 0

    def _create_extraction(self, *args) -> Extraction:
        return Extraction(*args)

    async def _extract_xml_content(self, data, collector: AbstractCollector, timeout: float = 60.0):
        async with self.session.put(
            f"http://127.0.0.1:{self.ec.tika_port}/tika", data=data, headers=self.TIKA_HEADERS, timeout=timeout
        ) as r:
            r.raise_for_status()
            pax = self._create_xml_parser(collector)
            try:
                async for data in r.content.iter_any():
                    pax.feed(data)
                pax.feed_end()
            except Exception as e:
                try:
                    pax.close()
                except etree.LxmlError:
                    pass  # we're already handling an error...
                if isinstance(e, ExtractError):
                    pass  # purposeful throw to interrupt parsing
                raise
            else:
                pax.close()

    def _create_xml_parser(self, collector) -> CollectorXMLParser:
        return CollectorXMLParser(collector)
