from __future__ import annotations

import asyncio
import copy
import math
import os
import random
import stat
import time
from contextlib import contextmanager, asynccontextmanager
from pathlib import Path, PureWindowsPath
from typing import Dict, Set, Iterable, List, Optional, Callable, AsyncContextManager

import wait_for2
from clarabot_nano_common.common.solr.rest import (
    nano_query_get,
    preparse_solr_response,
    extract_query_response,
    nano_update_post,
)
from clarabot_nano_common.indexer.utility import create_accessor_path
from clarabot_nano_common.utility.path.chroot import ChrootPath
from clarabot_nano_common.utility.path.common import normalize_absolute_path
from clarabot_nano_common.utility.path.resolve import resolve_path

from clarabot_nano.indexer.commit_tracker import DocumentCommitTracker
from clarabot_nano.indexer.config import IndexerConfig
from clarabot_nano.indexer.event import SchedulerFilesystemEvent
from clarabot_nano.indexer.initializer import IndexerSolrSchemaStateBindInitializerMixin
from clarabot_nano.indexer.scheduler import IndexingSchedulerEngine, SyncInterrupted
from clarabot_nano.indexer.sync_data import SyncFilesystemParent
from clarabot_nano_tests.testing.file import temp_path


class _PollSyncFinished(Exception):
    pass


class _EventSyncFinished(Exception):
    pass


class _IgnoreCommitTimeDocumentCommitTracker(DocumentCommitTracker):
    def add_commit_time(self, *args):
        pass


class _TestingIndexingSchedulerEngine(IndexingSchedulerEngine):

    # force testing query slicing by reduced window size
    DOC_SLICE_ITER_DIR = 5
    DOC_SLICE_DROP_TREE = 5

    FAILED_ROOTS_CHECK_INTERVAL = 0.25

    def __init__(self, *args, ignore_add_commit_time=False):
        IndexingSchedulerEngine.__init__(self, *args)
        if ignore_add_commit_time:
            # This allows ignoring soft-commits produced by elapsing time and commit-within.
            self._commit_tracker = _IgnoreCommitTimeDocumentCommitTracker(self._commit_tracker._max_time_until_commit)
        self._run_auto_soft_commits = True
        self._observe_fs_events = True
        self.observe_fs_started = asyncio.Event()
        self._poll_sync = True
        self._skip_event_aggregation = False
        self._simulate_interruptions = False
        self._interruption_count = 1
        self.soft_commit_count = 0
        self._sync_trigger_custom_count = 0
        self._custom_sync_callbacks: Dict[int, Callable] = {}
        self.exit_test_if_no_chroots = True
        self.observe_task_started = asyncio.Event()

    async def _poll_sync_traverse(self, chroots: Dict[str, ChrootPath]):
        if not self._poll_sync:
            while True:
                await asyncio.sleep(5)
        if not chroots and self.exit_test_if_no_chroots:
            # mark poll finished; so next test invocation will start right away
            await self.index_state_db.finish_poll_traversal(list(chroots.keys()))
            raise _PollSyncFinished()
        return await IndexingSchedulerEngine._poll_sync_traverse(self, chroots)

    async def _sync_path(self, chroot: ChrootPath, rel: str) -> Optional[SyncFilesystemParent]:
        custom_callback = self._custom_sync_callbacks.get(self._sync_trigger_custom_count)
        if custom_callback:
            custom_callback()
        self._sync_trigger_custom_count += 1
        return await IndexingSchedulerEngine._sync_path(self, chroot, rel)

    async def _process_event_aggregation_buckets(self):
        if self._skip_event_aggregation:
            while True:
                await self._event_aggregate_new_event.wait()
                self._event_aggregate_new_event.clear()
                while self._event_aggregate_buckets:
                    traversal_items = self._event_aggregate_buckets.popleft()[1]
                    await self.index_state_db.add_priority_traversal_items(traversal_items)
                    self._priority_traversal_added.set()
        else:
            await IndexingSchedulerEngine._process_event_aggregation_buckets(self)

    async def _solr_soft_commit(self):
        self.soft_commit_count += 1
        await IndexingSchedulerEngine._solr_soft_commit(self)

    async def _poll_sync_finish_delay(self, time_since_traversal_finish: float):
        if time_since_traversal_finish == 0.0:
            raise _PollSyncFinished()
        await IndexingSchedulerEngine._poll_sync_finish_delay(self, time_since_traversal_finish)

    async def _observe_drives(self, chroots: Dict[str, ChrootPath]):
        self.observe_task_started.set()
        if not self._observe_fs_events:
            while True:
                await asyncio.sleep(5)
        await IndexingSchedulerEngine._observe_drives(self, chroots)

    async def _observe_drives_loop(self, observer):
        self.observe_fs_started.set()
        await IndexingSchedulerEngine._observe_drives_loop(self, observer)

    async def _auto_soft_commit(self):
        if not self._run_auto_soft_commits:
            while True:
                await asyncio.sleep(5)
        await IndexingSchedulerEngine._auto_soft_commit(self)

    def _handle_processing_interrupt(self, context: int):
        # only simulate raising SyncInterrupted from polling (not from preparation or priority)
        if self._simulate_interruptions and self._sync_lock.locked() and self._sync_is_polling:
            if context == self.SYNC_INTERRUPT_CONTEXT_DELETE:
                raise SyncInterrupted()
            if random.random() < (1 / math.sqrt(self._interruption_count)):
                self._interruption_count += 1
                raise SyncInterrupted()
        IndexingSchedulerEngine._handle_processing_interrupt(self, context)

    async def run_single_poll_cycle(
        self,
        run_auto_soft_commits: bool = True,
        observe_fs_events: bool = False,
        simulate_interruptions: bool = True,
        skip_event_aggregation: bool = False,
        custom_sync_callbacks: Dict[int, Callable] = None,
        timeout: float = 20.0,
    ):
        self.observe_task_started.clear()
        self._run_auto_soft_commits = run_auto_soft_commits
        self._observe_fs_events = observe_fs_events
        if observe_fs_events:
            self.observe_fs_started.clear()
        else:
            self.observe_fs_started.set()
        self._poll_sync = True
        self._simulate_interruptions = simulate_interruptions
        self._skip_event_aggregation = skip_event_aggregation
        self._interruption_count = 1
        self.soft_commit_count = 0
        self._sync_trigger_custom_count = 0
        self._custom_sync_callbacks = {}
        if custom_sync_callbacks:
            self._custom_sync_callbacks.update(custom_sync_callbacks)
        try:
            await wait_for2.wait_for(self.run(), timeout)
        except _PollSyncFinished:
            pass
        else:
            assert False

    async def run_event_cycle(
        self,
        run_auto_soft_commits: bool = True,
        observe_fs_events: bool = False,
        simulate_interruptions: bool = True,
        skip_event_aggregation: bool = False,
    ):
        self.observe_task_started.clear()
        self._run_auto_soft_commits = run_auto_soft_commits
        self._observe_fs_events = observe_fs_events
        if observe_fs_events:
            self.observe_fs_started.clear()
        else:
            self.observe_fs_started.set()
        self._poll_sync = False
        self._simulate_interruptions = simulate_interruptions
        self._skip_event_aggregation = skip_event_aggregation
        self._interruption_count = 1
        self.soft_commit_count = 0
        self._sync_trigger_custom_count = 0
        self._custom_sync_callbacks = {}
        await self.run()


class AbstractDocsCollector(object):
    def __init__(self, session):
        self.session = session

    def get_solr_query_url(self):
        raise NotImplementedError()

    def get_solr_update_url(self):
        raise NotImplementedError()

    def get_solr_auth_header(self):
        raise NotImplementedError()

    def get_languages(self):
        raise NotImplementedError()

    def assert_accessor_in_roots(self, accessor_path: str):
        raise NotImplementedError()

    def get_assert_struct_from_docs(self, docs: Iterable[dict]):
        s = dict()
        for d in docs:
            assert d["path"] not in s
            accessors = set()
            for a in d["accessors"]:
                self.assert_accessor_in_roots(a)
                accessors.add(a)
            texts = []
            for key in d.keys():
                if key.startswith("text_") and d[key] is not None:
                    texts.append(d[key])
            if d["file"]:
                assert len(texts) == 1, str(texts)
            else:
                assert not texts
                texts.append(None)
            s[d["path"]] = {
                "file": d["file"],
                "accessors": accessors,
                "size": d.get("size"),
                "mtime": d.get("mtime") if d["file"] else None,  # omit volatile timestamp of directories
                "text": texts[0],
                "img_data": d.get("img_data"),
                "img_bucket": d.get("img_bucket"),
                "preview_data": d.get("preview_data"),
            }
        return s

    async def assert_from_paths(self, ps: Dict[str, dict]):
        async with nano_update_post(
            self.session, self.get_solr_update_url(), self.get_solr_auth_header(), {}, soft_commit=True
        ) as r:
            result, success = await preparse_solr_response(r)
            assert success, f"{r.status} {result}"
        lang_fields = ",".join(["text_faba"] + [f"text_{lang}" for lang in self.get_languages()])
        async with nano_query_get(
            self.session,
            self.get_solr_query_url(),
            self.get_solr_auth_header(),
            [
                ("defType", "lucene"),
                ("q", "*:*"),
                ("fl", f"path,file,accessors,size,mtime,{lang_fields},img_data,img_bucket,preview_data"),
                ("rows", n := 2000),
            ],
        ) as r:
            result, success = await preparse_solr_response(r)
            assert success, str(result)
        q = extract_query_response(result)
        assert (c := len(q["docs"])) < n, f"{c} {n}"
        ds = self.get_assert_struct_from_docs(q["docs"])
        fs_paths = set(ps.keys())
        doc_paths = set(ds.keys())
        assert fs_paths == doc_paths, (
            f"Docs mismatch {len(doc_paths)} {len(fs_paths)}"
            + f"\n  Missing from solr:\n    "
            + f"\n    ".join([f"{k}: {ps[k]}" for k in sorted(fs_paths.difference(doc_paths))])
            + f"\n  Missing from fs:\n    "
            + f"\n    ".join([f"{k}: {ds[k]}" for k in sorted(doc_paths.difference(fs_paths))])
        )
        for key in fs_paths:
            assert ps[key] == ds[key], f"File data mismatch for {key}\n  Doc {ds[key]}\n  Fs  {ps[key]}"
        assert ps == ds

    async def assert_from_paths_until(self, ps: Dict[str, dict], timeout=8.0):
        t = time.perf_counter() + timeout
        while True:
            try:
                return await self.assert_from_paths(ps)
            except AssertionError:
                self._check_wait_valid()
                if time.perf_counter() > t:
                    raise
                await asyncio.sleep(0.15)

    def _check_wait_valid(self):
        pass


class SchedulerTesterBase(AbstractDocsCollector):
    def __init__(self, session, config, index_state_db):
        AbstractDocsCollector.__init__(self, session)
        self.event_queue = asyncio.Queue()
        self._events_task = None
        self._original_config = config
        self._config = self.new_config()
        self.index_state_db = index_state_db
        self.ec = self._init_extracted_config(self.new_config())
        self.core_name = None
        self.languages = None
        self.core_config = None
        self.soft_commits = []
        self._old_roots = []

    @contextmanager
    def change_config(self, reset: bool = False):
        if reset:
            self._config = self.new_config()
        yield self._config
        self.ec = self._init_extracted_config(self._config)

    @contextmanager
    def old_roots(self, *roots: str):
        self._old_roots.extend(roots)
        try:
            yield
        finally:
            for r in roots:
                self._old_roots.remove(r)

    def generate_tmp_roots_change_config(self, tmpdir, num):
        rs = [Path(normalize_absolute_path(temp_path(tmpdir))).as_posix() for _ in range(num)]
        for r in rs:
            os.makedirs(r, exist_ok=True)
        with self.change_config() as config:
            config["remote"]["drives"] = [{"path": r} for r in rs]
        return rs

    def _init_extracted_config(self, config):
        ec = IndexerConfig(config)
        self.chroots, self.failed_roots = ec.get_drive_chroots()
        return ec

    # @property
    # def roots(self):
    #     return set(self.chroots)

    def get_path_roots(self, path: str):
        return set(root for root in self.chroots.keys() if path.startswith(root))

    def get_solr_query_url(self):
        return self.ec.solr_query_url

    def get_solr_update_url(self):
        return self.ec.solr_update_url

    def get_solr_auth_header(self):
        return self.ec.solr_auth_header

    def get_languages(self):
        return self.languages

    def assert_accessor_in_roots(self, accessor_path: str):
        if accessor_path.endswith("/"):
            root = Path(accessor_path).as_posix()
        else:
            root = accessor_path.rpartition("//")[0]
        # NOTE: this will fail for
        #   - if the failed root is a symlink
        assert (
            root in self.chroots or root in self.failed_roots or root in self._old_roots
        ), f"{root}\n{set(self.chroots.keys())}\n{self.failed_roots}\n{self._old_roots}"

    def _check_wait_valid(self):
        pass
        if self._events_task is not None and self._events_task.done():
            e = self._events_task.exception()
            if e:
                raise e
            raise Exception("events task stopped")

    def new_config(self):
        return copy.deepcopy(self._original_config)

    async def init_solr_core(self):
        (
            self.core_name,
            self.languages,
            self.core_config,
        ) = await IndexerSolrSchemaStateBindInitializerMixin()._initialise(self.session, self.ec, self.index_state_db)
        self._original_config["state"]["solr_core"] = self.core_name
        with self.change_config(True):
            pass  # reset with solr_core state

    def new_scheduler_engine(self, ignore_add_commit_time=False):
        return _TestingIndexingSchedulerEngine(
            self.session,
            self.ec,
            self.index_state_db,
            self.core_name,
            self.languages,
            self.core_config,
            self.event_queue,
            self._set_status,
            self._notify_work,
            ignore_add_commit_time=ignore_add_commit_time,
        )

    async def _set_status(self, status: str):
        pass

    def _notify_work(self, work: str):
        pass

    async def run_check_repeated_polls(
        self, psb: PathStructureBuilder, repeat: int = 3, **kwargs
    ) -> Optional[Dict[str, dict]]:
        ps = None
        for _ in range(repeat):
            e = self.new_scheduler_engine()
            await e.run_single_poll_cycle(**kwargs)
            self.soft_commits.append(e.soft_commit_count)
            ps = psb.generate_paths_structure()
            await self.assert_from_paths(ps)
        return ps

    @asynccontextmanager
    async def run_events(self, **kwargs) -> AsyncContextManager[_TestingIndexingSchedulerEngine]:
        assert self._events_task is None
        assert self.event_queue.empty()
        e = self.new_scheduler_engine()
        self._events_task = asyncio.create_task(e.run_event_cycle(**kwargs))
        await asyncio.wait(
            [
                asyncio.create_task(e.observe_task_started.wait()),  # generally when the engine is started
                asyncio.create_task(e.observe_fs_started.wait()),  # specifically when the fs observer is started
                asyncio.create_task(asyncio.sleep(0.5)),  # some minimal wait
            ],
            timeout=6.0,
            return_when=asyncio.ALL_COMPLETED,
        )
        await asyncio.sleep(0.1)  # wait a little more just to be sure?
        assert not self._events_task.done()
        try:
            yield e
        finally:
            self._events_task.cancel()
            await asyncio.gather(self._events_task, return_exceptions=True)
            self._events_task = None

    def create_event(self, root, *rel, exists=None):
        p = Path(root, *rel)
        if isinstance(p, PureWindowsPath) and random.random() > 0.5:
            p = p.as_posix()  # mix separators on windows for testing
        else:
            p = str(p)
        self.event_queue.put_nowait(SchedulerFilesystemEvent(p, exists))


class PathStructureBuilder(object):

    AUTO_FILE_NAME_CHARS = "xyz"

    def __init__(self, roots: Iterable[str]):
        self.roots = set(roots)
        self._lost_keeps = []
        self._file_name_id = 1
        self._content_id = 1

    def generate_content(self) -> str:
        c = f"Structure test file {self._content_id} {'0' * random.randint(1, 20)}"
        self._content_id += 1
        return c

    def generate_auto_name(self, chars: str, length: int) -> str:
        assert length > 0
        assert chars
        return "".join(random.choice(chars) for _ in range(length))

    def create_dirs(self, root: str, rels: Iterable[str]):
        for rel in rels:
            self.create_dir(root, rel)

    def create_dir(self, root: str, rel: str):
        Path(root, rel).mkdir(parents=True, exist_ok=True)

    def create_files_in_dirs(
        self, root: str, rel_parents: Dict[str, int], write_delay: float = 0.001
    ) -> Dict[str, List[str]]:
        rs = dict()
        for rel, count in rel_parents.items():
            rs[rel] = []
            for _ in range(count):
                name = f"{self.generate_auto_name(self.AUTO_FILE_NAME_CHARS, 3)}{self._file_name_id:03d}.txt"
                rs[rel].append(name)
                self._file_name_id += 1
                self.create_file(root, f"{rel}/{name}" if rel else name, write_delay=write_delay)
        return rs

    def delete_files_in_dirs(self, root: str, rel_parents: Dict[str, int]):
        for rel, count, files in self._random_files_from_dirs(root, rel_parents):
            for i in range(count):
                files[i].unlink()

    def update_files_in_dirs(self, root: str, rel_parents: Dict[str, int]):
        for rel, count, files in self._random_files_from_dirs(root, rel_parents):
            for i in range(count):
                files[i].write_text(self.generate_content())

    def _random_files_from_dirs(self, root, rel_parents):
        for rel, count in rel_parents.items():
            parent_dir = Path(root, rel)
            assert parent_dir.is_dir(), f"{root}\n{rel}"
            files = [p for p in parent_dir.iterdir() if stat.S_ISREG(p.lstat().st_mode)]
            assert len(files) >= count
            random.shuffle(files)
            yield rel, count, files

    def create_files(self, root: str, rels: Iterable[str], write_delay: float = 0.001):
        for rel in rels:
            self.create_file(root, rel, write_delay=write_delay)

    def create_file(self, root: str, rel: str, write_delay: float = 0.001):
        p = Path(root, rel)
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_text(self.generate_content())
        # NOTE: The length of the generated files often match, which means to test the behaviour of usual files,
        #   we need to ensure their timestamp is different. I.e. replacing sequentially generated file 1 with 0
        #   shall not fail the staleness check.
        time.sleep(write_delay)

    def generate_paths_structure(self) -> Dict[str, dict]:
        ps = self._generate_paths_structure()
        for lk in self._lost_keeps:
            ps.update(lk)
        return ps

    def _generate_paths_structure(self) -> Dict[str, dict]:
        ps = dict()
        for root in self.roots:
            self._walk_paths(ps, set(), resolve_path(root), root)
        return ps

    def _walk_paths(self, ps: Dict[str, dict], visited: Set[str], path: Path, root: str, *rel_parts: str):
        if not path.exists():
            return
        p = path.as_posix()
        if not p.startswith(root):
            return  # symlink containment simulation
        s = path.stat()  # resolve symlinks to target type
        if stat.S_ISREG(s.st_mode):
            assert rel_parts and all(rel_parts)
            self._walk_path_collect(
                ps,
                p,
                create_accessor_path(root, *rel_parts),
                True,
                s.st_size,
                s.st_mtime_ns,
                path.read_text(encoding="utf-8"),
            )
        else:
            if rel_parts:  # no directory record for root by default
                self._walk_path_collect(ps, p, create_accessor_path(root, *rel_parts))
            if p in visited:
                return  # avoid symlink cycles
            visited.add(p)
            items = list(path.iterdir())
            if not rel_parts:
                # directory record for root is added by their direct child files if any
                for item in items:
                    try:
                        if stat.S_ISREG(item.stat().st_mode):
                            self._walk_path_collect(ps, p, create_accessor_path(root))
                            break
                    except FileNotFoundError:
                        pass
            for item in items:
                self._walk_paths(ps, visited, resolve_path(item), root, *rel_parts, item.name)
            visited.remove(p)

    def _walk_path_collect(
        self, ps: Dict[str, dict], path: str, accessor_path: str, file=False, size=None, mtime=None, text=None
    ):
        name = path.rpartition("/")[-1]
        if file and name.startswith(".~CbN.") and name.endswith("#"):
            return
        if path not in ps:
            ps[path] = doc = {"accessors": set()}
        else:
            doc = ps[path]
        doc["file"] = file
        doc["accessors"].add(accessor_path)
        doc["size"] = size
        doc["mtime"] = mtime
        doc["text"] = text
        doc["img_data"] = None
        doc["img_bucket"] = None
        doc["preview_data"] = None

    def generate_filtered_paths_structure(self, f: Dict[str, List[Iterable[str]]]) -> Dict[str, dict]:
        ps = self.generate_paths_structure()
        fps = dict()
        for root, path_part_list in f.items():
            rp = Path(root)
            rpp = rp.as_posix()
            for path_parts in path_part_list:
                if not path_parts:  # include directory doc of root
                    if rpp not in fps:
                        assert rpp in ps, f"{rpp}\n{ps}"
                        fps[rpp] = ps[rpp]
                    continue
                p = Path(root, *path_parts)
                pp = p.as_posix()
                assert pp in ps, f"{pp}\n{ps}"
                fps[pp] = ps[pp]
                for i in range(len(rp.parts) + 1, len(p.parts)):
                    pp = Path(*p.parts[:i]).as_posix()
                    assert pp in ps, f"{pp}\n{ps}"
                    fps[pp] = ps[pp]
                if p.is_file() and p.parent == rp and rpp not in fps:  # add directory of root if not added already
                    assert rpp in ps, f"{rpp}\n{ps}"
                    fps[rpp] = ps[rpp]
        return fps

    @contextmanager
    def keep_paths_from_lost(self, root, rels):
        ps = self._generate_paths_structure()
        keeps = {Path(root, rel).as_posix() for rel in rels}
        if not keeps:
            keeps.add(Path(root).as_posix())
        kl = dict(
            (k, v)
            for k, v in ps.items()
            if any(k.startswith(p) and (len(k) == len(p) or k[len(p)] == "/") for p in keeps)
        )
        self._lost_keeps.append(kl)
        try:
            yield
        finally:
            self._lost_keeps.remove(kl)
