from __future__ import annotations

import asyncio
import itertools
import logging
from typing import Optional, List

import aiohttp
from lxml import etree

from clarabot_nano.indexer.extractor.base import (
    AbstractExtraction,
    SPACE_SEARCH_RE,
    NEWLINE_STRIP_RE,
    AbstractContentExtractor,
    ExtractError,
    AbstractExtractionSource,
    ExtractFailed,
)
from clarabot_nano.indexer.extractor.xml.base import AbstractCollector, CollectorHandler
from clarabot_nano.indexer.extractor.xml.body import BodyHandler
from clarabot_nano.indexer.extractor.xml.head import TitleHandler, HeadHandler

logger = logging.getLogger(__name__)


class TikaExtraction(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 TikaContentExtractor(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",
        # "X-Tika-PDFOcrStrategy": "NO_OCR",
        # "X-Tika-PDFMaxMainMemoryBytes": "300000000",
    }

    def _create_extraction(self, source: AbstractExtractionSource) -> TikaExtraction:
        return TikaExtraction(
            self.ec.indexer_ignored_mimes, self.ec.indexer_extract_text_limit, source.fallback_summary()
        )

    def _is_status_successful(self, source: AbstractExtractionSource, ex: AbstractExtraction, status: int) -> bool:
        """
        Some 4xx statuses may be considered successful for some content files.
        """
        return status == 422  # 422 is returned for empty files; so it's a success...

    async def check_operable(self, timeout: float = 5.0) -> bool:
        ex = TikaExtraction(set(), 65536, "fallback")
        try:
            await self._extract_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

    async def _extract_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:
                try:
                    async for data in r.content.iter_any():
                        pax.feed(data)
                    pax.feed_end()
                except Exception:
                    try:
                        pax.close()
                    except etree.LxmlError:
                        pass  # we're already handling an error...
                    raise
                else:
                    pax.close()
            except etree.LxmlError as e:
                # Handle unexpected errors. For example: a "Numbers" file has produced an empty Tika response before,
                # which means there was some internal parsing issue that prevented the response from being generated,
                # making the output invalid for the XML parser.
                raise ExtractFailed() from e  # Retry, then abandon if issue persists.

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