import asyncio
import json
import os
import sys
from typing import List, NamedTuple

import pytest

from clarabot_nano.indexer.extractor.base import (
    ALL_SPACE_SEARCH_RE,
    SPACE_SEARCH_RE,
    NEWLINE_SEARCH_RE,
    NEWLINE_STRIP_RE,
    ServerInternalExtractError,
    RetryExtractError,
    ExtractorInoperable,
)
from clarabot_nano_tests.indexer.extractor.resources import INPUT_ROOT, EXTRACT_ROOT, OUTPUT_ROOT, SAMPLES_ROOT
from clarabot_nano_tests.indexer.extractor.utility import ExtractorTesterBase, DiffXMLTrees
from clarabot_nano_tests.picoture.resources import INPUT_ROOT as PICOTURE_INPUT_ROOT

HEAD_INDEX = 1


def _move_head_node_around(x: DiffXMLTrees, d: List[NamedTuple]) -> int:
    if d[0].__class__.__name__ != "MoveNode":
        return 0
    if not x.paths_match([None, HEAD_INDEX, -1], x.list_path(d[0].node)):
        return 0
    if not x.paths_match([None, HEAD_INDEX], x.list_path(d[0].target)):
        return 0
    return 1


def _volatile_metadata_content(x: DiffXMLTrees, d: List[NamedTuple]) -> int:
    if d[0].__class__.__name__ != "UpdateAttrib" or d[0].name != "content":
        return 0
    l, r = x.nodes_from_path(d[0].node)
    if len(l) != 1 or len(r) != 1:
        return 0
    if not l[0].tag.endswith("}meta") or not r[0].tag.endswith("}meta"):
        return 0
    for name in ["File Name", "File Modified Date"]:
        if l[0].attrib.get("name") == name and r[0].attrib.get("name") == name:
            return 1
    return 0


def _extra_meta_content_type_override_win(x: DiffXMLTrees, d: List[NamedTuple]) -> int:
    if d[0].__class__.__name__ != "InsertNode" or not d[0].tag.endswith("}meta"):
        return 0
    p0 = x.list_path(d[0].target)
    if not x.paths_match([None, HEAD_INDEX], p0):
        return 0
    attr_path = [*p0, d[0].position + 1]
    attrs = {}
    for i in range(1, 3):
        if d[i].__class__.__name__ != "InsertAttrib" or not x.paths_match(attr_path, x.list_path(d[i].node)):
            return 0
        attrs[d[i].name] = d[i].value
    if (
        set(attrs.keys()) != {"name", "content"}
        or attrs["name"] != "Content-Type-Override"
        or not (
            attrs["content"].startswith("application/vnd.openxmlformats-officedocument.")
            or attrs["content"].startswith("application/vnd.ms-")
        )
    ):
        return 0
    return 3


EXTRACT_DIFFERENCE_ACCEPTORS = [_move_head_node_around, _volatile_metadata_content]

if sys.platform.startswith("win"):
    EXTRACT_DIFFERENCE_ACCEPTORS.append(_extra_meta_content_type_override_win)


def test_indexer_extractor_regexps():
    assert ALL_SPACE_SEARCH_RE.sub("", "  \xa0 a b \n c\nd  \r\n\t\r  \n  e ") == "abcde"
    assert SPACE_SEARCH_RE.sub("", "  a b \n c\nd  e ") == "ab\nc\nde"
    assert SPACE_SEARCH_RE.sub("", "  \xa0 a   b \n   c\nd  \xa0\t e\tx ") == "ab\nc\ndex"
    assert NEWLINE_SEARCH_RE.sub("", "a\nb\r\nc\rd\r\re\n\nf \r\n\r\r\n\r\n\n\n\r\r\ng") == "abcdef g"
    assert NEWLINE_STRIP_RE.sub("\n", "\r\n \t\n a\n \t b \r\n\t\r  \n c \n  \t") == "\na\nb\nc\n"


async def _check_extract_output(indexer_extractor_tester: ExtractorTesterBase, ip, ep, op):
    extractor = indexer_extractor_tester.create_extractor()
    # Limit is reduced to approx because of final normalization.
    approx_length_limit = indexer_extractor_tester.ec.indexer_extract_text_limit * 0.9
    for in_path, ex_dir, out_dir in indexer_extractor_tester.walk_resources(ip, ep, op):
        # for debugging
        # if in_path.name != "tf33478217_win32.dotx":
        #     continue
        assert in_path.is_file(), f"{in_path}"
        attempt = 0
        async with indexer_extractor_tester.create_file_extraction_source(in_path).prepare() as fes:
            extract_info = json.dumps(
                {
                    "content_type": fes.content_type,
                    "extract_text": fes.extract_text,
                    "extract_image": fes.extract_image,
                    "buffer_content": isinstance(fes.content, bytearray),
                },
                indent=4,
            ).encode("utf-8")
            while True:
                try:
                    extraction = await extractor.extract(fes)
                except ServerInternalExtractError as e:
                    extraction = e.extraction
                except ExtractorInoperable:
                    assert False, f"{in_path}"
                except RetryExtractError:  # OOM may crash the worker
                    attempt += 1
                    if attempt <= 2:
                        await asyncio.sleep(0.2)
                    else:
                        raise
                    continue
                break
        # verify/write extracts
        extracts = [
            extract.replace(b"\r\n", b"\n") for extract in [extract_info, *extraction.get_extracts()]
        ]  # fix newlines
        # for debugging
        # info_json_path = ex_dir / "0.json"
        # if info_json_path.is_file():
        #     info_json_path.write_bytes(extracts[0])
        extracts_in_dir = indexer_extractor_tester.read_dir_files(ex_dir)
        if extracts_in_dir is not None:
            assert len(extracts) == len(extracts_in_dir), f"{in_path}"
            for i, (extract, extract_in_dir) in enumerate(zip(extracts, extracts_in_dir)):
                if i == 0:
                    assert extract == extract_in_dir, (
                        f"extract infos mismatch {in_path}:\n  "
                        f"generated: {extract}\n  reference: {extract_in_dir}\n  =="
                    )
                else:
                    ex_diff = DiffXMLTrees.create(extract, extract_in_dir)
                    if not ex_diff.is_difference_acceptable(EXTRACT_DIFFERENCE_ACCEPTORS):
                        raise AssertionError(
                            f"extracts mismatch {i} {in_path}:\n  "
                            f"generated: {extract}\n  reference: {extract_in_dir}\n  {ex_diff}\n  =="
                        )
        elif extracts:
            ex_dir.mkdir(parents=True)
            for i, extract in enumerate(extracts):
                (ex_dir / f"{i}.{'json' if i == 0 else 'xml'}").write_bytes(extract)
        # verify/write outputs
        outputs = extraction.get_outputs()
        outputs_in_dir = indexer_extractor_tester.read_dir_files(out_dir, raw=False)
        if outputs_in_dir is not None:
            assert len(outputs) == len(outputs_in_dir), f"{in_path}"
            for generated, reference in zip(outputs, outputs_in_dir):
                if len(generated) != len(reference):
                    # Tika extractor yields fragments in different increments on different platforms and maybe configs.
                    # So the end of an extract cut may be different, but match up to the existing values.
                    shorter_length = min(len(generated), len(reference))
                    assert approx_length_limit < shorter_length, f"{in_path}\n{approx_length_limit}\n{shorter_length}"
                    assert generated[:shorter_length] == reference[:shorter_length], f"{in_path}"
                else:
                    assert generated == reference, f"{in_path}"
        else:
            out_dir.mkdir(parents=True)
            for i, output in enumerate(outputs):
                (out_dir / f"{i}.txt").write_text(output, encoding="utf-8")


# NOTE: System locale affects date, currency and some number formatting. The references are made with English/US.
@pytest.mark.asyncio
async def test_indexer_extractor_included(indexer_extractor_tester: ExtractorTesterBase):
    await _check_extract_output(indexer_extractor_tester, INPUT_ROOT, EXTRACT_ROOT, OUTPUT_ROOT)


@pytest.mark.asyncio
async def test_indexer_extractor_picoture(indexer_extractor_tester: ExtractorTesterBase):
    _dir = os.path.dirname(EXTRACT_ROOT)
    await _check_extract_output(
        indexer_extractor_tester,
        PICOTURE_INPUT_ROOT,
        os.path.join(_dir, "picoture_extract"),
        os.path.join(_dir, "picoture_output"),
    )


@pytest.mark.skipif(not os.path.exists(SAMPLES_ROOT), reason="can't find nano-sample-files repo")
# @pytest.mark.parametrize("run", list(range(50)))
@pytest.mark.asyncio
async def test_indexer_extractor_external(indexer_extractor_tester: ExtractorTesterBase):
    await _check_extract_output(
        indexer_extractor_tester,
        os.path.join(SAMPLES_ROOT, "input"),
        os.path.join(SAMPLES_ROOT, "extract"),
        os.path.join(SAMPLES_ROOT, "output"),
    )
