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

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
from clarabot_nano.indexer.extractor.file import Extraction, CollectorXMLParser, FileContentExtractor
from clarabot_nano_tests.testing.core.instance import CoreTestingNanoInstanceHandler
from clarabot_nano_tests.testing.tika import wait_tika_startup


class _TestingExtraction(Extraction):
    def __init__(self, ignored_mimes: set, extract_text_limit: int, fallback_summary: str):
        self.extracts = []
        Extraction.__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: _TestingExtraction):
        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):
        self.ex.extracts.append(self.raw_data.getvalue())
        self.raw_data.close()
        CollectorXMLParser.close(self)


class _TestingFileContentExtractor(FileContentExtractor):
    async def extract(self, path: Path) -> _TestingExtraction:  # fix type hinting
        return await FileContentExtractor.extract(self, path)

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

    def _create_xml_parser(self, collector):
        assert isinstance(collector, _TestingExtraction)
        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 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))


def diff_extracts(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 xmldiff.main.diff_trees(left_tree, right_tree, diff_options=diff_options, formatter=formatter)
