from __future__ import annotations

import copy
import os
from contextlib import asynccontextmanager
from io import BytesIO
from pathlib import Path
from typing import Generator, Tuple, List, Optional, Union, Callable, NamedTuple

import xmldiff.formatting
import xmldiff.main
from lxml import etree

from clarabot_nano.indexer.config import IndexerConfig
from clarabot_nano.indexer.extractor.base import ExtractError, AbstractExtractionSource
from clarabot_nano.indexer.extractor.source import FileExtractionSource
from clarabot_nano.indexer.extractor.tika import TikaExtraction, CollectorXMLParser, TikaContentExtractor
from clarabot_nano.indexer.extractor.xml.base import AbstractCollector
from clarabot_nano_tests.testing.core.instance import CoreTestingNanoInstanceHandler
from clarabot_nano_tests.testing.tika import wait_tika_startup


class _TestingTikaExtraction(TikaExtraction):
    def __init__(self, ignored_mimes: set, extract_text_limit: int, fallback_summary: str):
        self.extracts = []
        TikaExtraction.__init__(self, ignored_mimes, extract_text_limit, fallback_summary)

    def get_extracts(self) -> List[bytes]:
        return self.extracts

    def get_outputs(self) -> List[str]:
        return [self.get_summary(), self.get_content()]


class _TestingCollectorXMLParser(CollectorXMLParser):
    def __init__(self, collector: AbstractCollector):
        CollectorXMLParser.__init__(self, collector)
        self.ex = collector
        self.raw_data = BytesIO()
        self._extract_error = None

    def feed(self, data):
        self.raw_data.write(data)
        if self._extract_error is None:
            try:
                CollectorXMLParser.feed(self, data)
            except ExtractError as e:
                self._extract_error = e

    def feed_end(self):
        if self._extract_error is not None:
            raise self._extract_error
        else:
            CollectorXMLParser.feed_end(self)

    def close(self):
        if isinstance(self.ex, _TestingTikaExtraction):
            self.ex.extracts.append(self.raw_data.getvalue())
        self.raw_data.close()
        CollectorXMLParser.close(self)


class _TestingFileContentExtractor(TikaContentExtractor):
    async def extract(self, source: AbstractExtractionSource) -> _TestingTikaExtraction:  # fix type hinting
        return await TikaContentExtractor.extract(self, source)

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

    def _create_xml_parser(self, collector):
        return _TestingCollectorXMLParser(collector)


class ExtractorTesterBase(object):
    def __init__(self, session, config):
        self.session = session
        self.config = config
        self.ec = IndexerConfig(self.config)

    def create_extractor(self) -> _TestingFileContentExtractor:
        return _TestingFileContentExtractor(self.session, self.ec)

    def create_file_extraction_source(self, path: Path) -> FileExtractionSource:
        return FileExtractionSource(self.session, self.ec, path.name, path, path.stat())

    def walk_resources(
        self, input_root: Union[str, Path], extract_root: Union[str, Path], output_root: Union[str, Path]
    ) -> Generator[Tuple[Path, Path, Path], None, None]:
        input_root = Path(input_root)
        extract_root = Path(extract_root)
        output_root = Path(output_root)
        input_paths = set()
        for root, _, files in os.walk(input_root):
            for file in files:
                input_paths.add(Path(root, file).relative_to(input_root))
        for root, _, files in os.walk(extract_root):
            for file in files:
                rdp = Path(root, file).parent.relative_to(extract_root)
                assert rdp in input_paths
        for root, _, files in os.walk(output_root):
            for file in files:
                rdp = Path(root, file).parent.relative_to(output_root)
                assert rdp in input_paths
        for rfp in input_paths:
            yield input_root / rfp, extract_root / rfp, output_root / rfp

    def read_dir_files(self, dir_path: Path, raw=True) -> Optional[Union[List[str], List[bytes]]]:
        if not dir_path.exists():
            return
        ps = []
        for p in dir_path.iterdir():
            i = p.name.rpartition(".")[0]
            assert i.isnumeric()
            ps.append((int(i), p))
        ps.sort(key=lambda x: x[0])
        assert all(px[0] == i for px, i in zip(ps, range(len(ps))))
        if raw:
            return [px[1].read_bytes() for px in ps]
        return [px[1].read_text("utf-8") for px in ps]


@asynccontextmanager
async def start_indexer_extractor_test_base(reset_user_app_dirs, dev_nano_root_dir, tika_heap_size=None):
    def _tika_heap(conf):
        conf["tika_heap_size"] = tika_heap_size

    async with CoreTestingNanoInstanceHandler(
        reset_user_app_dirs, dev_nano_root_dir, local_config_mut=_tika_heap if tika_heap_size else None
    ).run() as (
        pm,
        session,
        children_spawns,
        dbg_config,
    ):
        ec_ = IndexerConfig(dbg_config)
        await wait_tika_startup(session, ec_.tika_port)
        yield ExtractorTesterBase(session, copy.deepcopy(dbg_config))


class DiffXMLTrees(object):
    @classmethod
    def create(cls, a: bytes, b: bytes, diff_options=None, formatter=None):
        normalize = bool(getattr(formatter, "normalize", 1) & xmldiff.formatting.WS_TAGS)
        parser = etree.XMLParser(
            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=normalize,
            resolve_entities=True,
            remove_comments=False,
            remove_pis=False,
            strip_cdata=False,
            collect_ids=False,
            compact=False,
        )
        left_tree = etree.parse(BytesIO(a), parser)
        right_tree = etree.parse(BytesIO(b), parser)
        return cls(left_tree, right_tree, diff_options=diff_options, formatter=formatter)

    def __init__(self, left: etree.ElementTree, right: etree.ElementTree, diff_options=None, formatter=None):
        self.left = left
        self.right = right
        self.diff = xmldiff.main.diff_trees(left, right, diff_options=diff_options, formatter=formatter)

    def is_difference_acceptable(self, acceptors: List[Callable[[DiffXMLTrees, List[NamedTuple]], int]]):
        d = list(self.diff)
        while d:
            accepted = 0
            for acceptor in acceptors:
                accepted = acceptor(self, d)
                if accepted > 0:
                    d = d[accepted:]
                    break
            if accepted == 0:
                return False
        return True

    def nodes_from_path(self, t: str) -> Tuple[List[etree.Element], List[etree.Element]]:
        return self.left.xpath(t), self.right.xpath(t)

    def left_node_from_path(self, t: str) -> Optional[etree.Element]:
        x = self.left.xpath(t)
        if len(x) == 1:
            return x[0]

    def right_node_from_path(self, t: str) -> Optional[etree.Element]:
        x = self.right.xpath(t)
        if len(x) == 1:
            return x[0]

    @staticmethod
    def list_path(t: str) -> List[Optional[int]]:
        p = []
        for x in t.split("/"):
            if not x:
                pass
            elif x == "*":
                p.append(None)
            elif x.startswith("*[") and x.endswith("]"):
                p.append(int(x[2:-1]))
            else:
                assert False, t
        return p

    @staticmethod
    def paths_match(test_path: List[Optional[int]], check_path: List[Optional[int]]) -> bool:
        if len(check_path) < len(test_path):
            return False
        for tp, cp in zip(test_path, check_path):
            if tp is None:
                if cp is not None:
                    return False
            elif tp == -1:
                if cp is None:
                    return False
            elif tp != cp:
                return False
        return True
