import asyncio
import base64
import itertools
import sys
from collections import defaultdict
from functools import partial

import pytest

from clarabot_nano.common.solr.rest import (
    create_nano_query_url,
    create_nano_update_url,
    create_solr_core_url,
    check_solr_error,
    create_solr_auth_header,
    serialize_solr_update_json,
)
from clarabot_nano.indexer.utility import (
    create_accessor_path,
    create_docs_query_tree,
    create_docs_query_exact,
    create_docs_query_children,
)
from clarabot_nano.utility import parse_environment_flag
from clarabot_nano_tests.testing.core.instance import ConfigTestingNanoInstanceHandler
from clarabot_nano_tests.testing.solr import (
    run_check_query,
    run_check_update,
    SOLR_ALL_LANGS,
    mutate_remote_config_indexed_languages,
)


@pytest.mark.skipif(parse_environment_flag("CBNANO_TEST_WITHOUT_SERVER"), reason="server-dependant tests disabled")
@pytest.mark.asyncio
async def test_indexer_solr_schema(testing_server_manager, reset_user_app_dirs, dev_nano_root_dir):
    async with ConfigTestingNanoInstanceHandler(
        reset_user_app_dirs,
        dev_nano_root_dir,
        testing_server_manager,
        partial(mutate_remote_config_indexed_languages, "all"),
    ).run() as (pm, debug_info, session, dbg_config):
        solr_port = dbg_config["state"]["solr_port"]
        solr_core = dbg_config["state"]["solr_core"]
        assert create_solr_core_url("127.0.0.2", solr_port, "xy", "zq") == f"http://127.0.0.2:{solr_port}/solr/xy/zq"
        q_url = create_nano_query_url(solr_port, solr_core)
        u_url = create_nano_update_url(solr_port, solr_core)
        auth_h = create_solr_auth_header(dbg_config["state"]["solr_key"])

        async with session.get(
            f"http://127.0.0.1:{solr_port}/solr/admin/cores?action=STATUS&core={solr_core}",
            headers={"Authorization": auth_h},
        ) as r:
            d = await r.json()
            assert d["responseHeader"]["status"] == 0, f"{d}"
            assert d["status"][solr_core]["name"] == solr_core, f"{d}"

        def _add_lang_fileds(doc_):
            for lang in SOLR_ALL_LANGS:
                doc_[f"summary_{lang}"] = doc_["summary_faba"]
                doc_[f"text_{lang}"] = doc_["text_faba"]

        # check basic INSERT & UPDATE & DELETE
        q0_args = [
            ("q", 'path:"/tmp/xy"'),
            ("fl", "path,file,accessors,size,mtime"),
            ("start", 0),
            ("rows", 2),
            ("sort", "mtime desc"),
        ]
        q1_args = list(q0_args)
        q1_args[0] = ("q", 'path:"/tmp/asd"')
        for q_args in [q0_args, q1_args]:
            await run_check_query(session, q_url, auth_h, q_args, [], 0)
        # add first doc
        docs0 = [
            {
                "path": "/tmp/xy",
                "file": True,
                "accessors": [create_accessor_path("/tmp", "xy")],
                "size": 77,
                "mtime": 555,
            }
        ]
        await run_check_update(session, u_url, auth_h, docs0)
        # re-check queries
        await run_check_query(session, q_url, auth_h, q0_args, docs0, 1)
        await run_check_query(session, q_url, auth_h, q1_args, [], 0)
        # update + add second doc
        docs1 = [
            {
                "path": "/tmp/xy",
                "file": True,
                "accessors": [create_accessor_path("/tmp", "xy")],
                "size": 50,
                "mtime": 8000,
            },
            {
                "path": "/tmp/asd",
                "file": True,
                "accessors": [create_accessor_path("/tmp", "asd")],
                "size": 50,
                "mtime": 8000,
            },
        ]
        await run_check_update(session, u_url, auth_h, docs1)
        # re-check queries
        for q_args, doc in zip([q0_args, q1_args], docs1):
            await run_check_query(session, q_url, auth_h, q_args, [doc], 1)
        # delete updated doc
        await run_check_update(session, u_url, auth_h, {"delete": {"id": "/tmp/xy"}})
        # re-check queries
        await run_check_query(session, q_url, auth_h, q0_args, [], 0)
        await run_check_query(session, q_url, auth_h, q1_args, [docs1[1]], 1)

        # check partial/atomic UPDATE behaviour on all fields (by accessors)
        lang_fields = []
        for lang in SOLR_ALL_LANGS:
            lang_fields.append(f"summary_{lang}")
            lang_fields.append(f"text_{lang}")
        pq0_args = [
            ("q", f'accessors:"{create_accessor_path("/tmp/xy")}"'),
            (
                "fl",
                f"path,file,accessors,size,mtime,summary_faba,text_faba,{','.join(lang_fields)},img_data,img_bucket",
            ),
            ("start", 0),
            ("rows", 10),
            ("sort", "mtime desc"),
        ]
        pq1_args = list(pq0_args)
        pq1_args[0] = ("q", f'accessors:"{create_accessor_path("/var/x")}"')
        await run_check_query(session, q_url, auth_h, pq0_args, [], 0)
        await run_check_query(session, q_url, auth_h, pq1_args, [], 0)
        d0 = {
            "path": "/tmp/xy/asd/x0",
            "file": True,
            "accessors": {"add-distinct": create_accessor_path("/tmp/xy", "asd/x0")},
            "size": 80,
            "mtime": 601,
            "summary_faba": "Title of X0",
            "text_faba": "X0 is a long story of a person who would not want to grasp reality...",
            "img_data": base64.b64encode(b"xa").decode("ascii"),
            "img_bucket": 1,
        }
        _add_lang_fileds(d0)
        await run_check_update(session, u_url, auth_h, [d0])
        d0["accessors"] = [d0["accessors"]["add-distinct"]]
        await run_check_query(session, q_url, auth_h, pq0_args, [d0], 1)
        await run_check_query(session, q_url, auth_h, pq1_args, [], 0)
        d1 = {
            "path": "/tmp/xy/asd/x1",
            "file": True,
            "accessors": {"add-distinct": create_accessor_path("/tmp/xy", "asd/x1")},
            "size": 709,
            "mtime": 707,
            "summary_faba": "Title of X1",
            "text_faba": "X1 is the story of the end... It came when everyone knew yet none believed.",
            "img_data": base64.b64encode(b"6V").decode("ascii"),
            "img_bucket": 1,
        }
        _add_lang_fileds(d1)
        await run_check_update(session, u_url, auth_h, [d1])
        d1["accessors"] = [d1["accessors"]["add-distinct"]]
        await run_check_query(session, q_url, auth_h, pq0_args, [d1, d0], 2)
        await run_check_query(session, q_url, auth_h, pq1_args, [], 0)
        d2 = {
            "path": "/tmp/xy/x2",
            "file": True,
            "accessors": {"add-distinct": create_accessor_path("/tmp/xy", "x2")},
            "size": 709,
            "mtime": 657,
            "summary_faba": "Title of X2",
            "text_faba": "X2 describes how nothingness enveloped everything...",
            "img_data": base64.b64encode(b"WQ").decode("ascii"),
            "img_bucket": 1,
        }
        _add_lang_fileds(d2)
        await run_check_update(session, u_url, auth_h, [d2])
        d2["accessors"] = [d2["accessors"]["add-distinct"]]
        await run_check_query(session, q_url, auth_h, pq0_args, [d1, d2, d0], 3)
        await run_check_query(session, q_url, auth_h, pq1_args, [], 0)
        # test add symlink path with different root
        await run_check_update(
            session,
            u_url,
            auth_h,
            [
                {
                    "path": d2["path"],
                    "accessors": {"add-distinct": (a_path := create_accessor_path("/var/x", "symlink/x2"))},
                }
            ],
        )
        d2["accessors"].append(a_path)
        await run_check_query(session, q_url, auth_h, pq0_args, [d1, d2, d0], 3)
        await run_check_query(session, q_url, auth_h, pq1_args, [d2], 1)
        # test add symlink path with same root
        await run_check_update(
            session,
            u_url,
            auth_h,
            [
                {
                    "path": d2["path"],
                    "accessors": {"add-distinct": (a_path := create_accessor_path("/tmp/xy", "alt/x2sl"))},
                }
            ],
        )
        d2["accessors"].append(a_path)
        await run_check_query(session, q_url, auth_h, pq0_args, [d1, d2, d0], 3)
        await run_check_query(session, q_url, auth_h, pq1_args, [d2], 1)
        # test update normal fields and retain accessors
        d2_update = {
            "path": d2["path"],
            "file": True,
            "accessors": {"add-distinct": d2["accessors"][0]},
            "size": 805,
            "mtime": 1200,
            "summary_faba": d2["summary_faba"] + " - rev 2",
            "text_faba": d2["text_faba"] + "\nsecond revision",
            "img_data": base64.b64encode(b"K2").decode("ascii"),
            "img_bucket": 1,
        }
        _add_lang_fileds(d2_update)
        await run_check_update(session, u_url, auth_h, [d2_update])
        d2_update.pop("accessors")
        d2.update(d2_update)
        await run_check_query(session, q_url, auth_h, pq0_args, [d2, d1, d0], 3)
        await run_check_query(session, q_url, auth_h, pq1_args, [d2], 1)
        # test accessor removals
        await run_check_update(
            session,
            u_url,
            auth_h,
            [{"path": d2["path"], "accessors": {"remove": (r_path := create_accessor_path("/tmp/xy", "x2"))}}],
        )
        d2["accessors"].remove(r_path)
        await run_check_query(session, q_url, auth_h, pq0_args, [d2, d1, d0], 3)
        await run_check_query(session, q_url, auth_h, pq1_args, [d2], 1)
        await run_check_update(
            session,
            u_url,
            auth_h,
            [{"path": d2["path"], "accessors": {"remove": (r_path := create_accessor_path("/tmp/xy", "alt/x2sl"))}}],
        )
        d2["accessors"].remove(r_path)
        await run_check_query(session, q_url, auth_h, pq0_args, [d1, d0], 2)
        await run_check_query(session, q_url, auth_h, pq1_args, [d2], 1)
        upd = await run_check_update(
            session,
            u_url,
            auth_h,
            [{"path": d2["path"], "accessors": {"remove": (f_path := create_accessor_path("/var/x", "symlink/x2"))}}],
            success=False,
        )
        check_solr_error(upd, 400, "missing required field: accessors")
        await run_check_update(session, u_url, auth_h, {})  # commit
        await run_check_query(session, q_url, auth_h, pq0_args, [d1, d0], 2)
        await run_check_query(session, q_url, auth_h, pq1_args, [d2], 1)
        # try multi-add to accessors
        await run_check_update(
            session,
            u_url,
            auth_h,
            [
                {
                    "path": d2["path"],
                    "accessors": {"add-distinct": (a_path := create_accessor_path("/var/x", "alt/qq0"))},
                },
                {
                    "path": d2["path"],
                    "accessors": {"add-distinct": (b_path := create_accessor_path("/var/x", "alt/qq1"))},
                },
                {
                    "path": d2["path"],
                    "accessors": {"add-distinct": (c_path := create_accessor_path("/var/x", "alt/qq2"))},
                },
            ],
        )
        d2["accessors"].extend([a_path, b_path, c_path])
        await run_check_query(session, q_url, auth_h, pq0_args, [d1, d0], 2)
        await run_check_query(session, q_url, auth_h, pq1_args, [d2], 1)
        # try multi-remove from accessors
        await run_check_update(
            session,
            u_url,
            auth_h,
            [
                {"path": d2["path"], "accessors": {"remove": a_path}},
                {"path": d2["path"], "accessors": {"remove": c_path}},
            ],
        )
        d2["accessors"].remove(a_path)
        d2["accessors"].remove(c_path)
        await run_check_query(session, q_url, auth_h, pq0_args, [d1, d0], 2)
        await run_check_query(session, q_url, auth_h, pq1_args, [d2], 1)
        # try multi-remove from accessors
        upd = await run_check_update(
            session,
            u_url,
            auth_h,
            [
                {"path": d2["path"], "accessors": {"remove": b_path}},
                {"path": d2["path"], "accessors": {"remove": f_path}},
            ],
            success=False,
        )
        check_solr_error(upd, 400, "missing required field: accessors")  # delay commit
        await run_check_query(session, q_url, auth_h, pq0_args, [d1, d0], 2)
        await run_check_query(session, q_url, auth_h, pq1_args, [d2], 1)  # not committed, but only has 1 accessor left
        await run_check_update(session, u_url, auth_h, {})  # commit
        d2["accessors"].remove(b_path)
        await run_check_query(session, q_url, auth_h, pq1_args, [d2], 1)
        # remove
        await run_check_update(session, u_url, auth_h, {"delete": {"id": d2["path"]}})
        await run_check_query(session, q_url, auth_h, pq0_args, [d1, d0], 2)
        await run_check_query(session, q_url, auth_h, pq1_args, [], 0)

        # == strict accessor root matching tests ==
        # The goal of the accessors tokenization is to provide a path-hierarchy (parent) lookup while also including
        # information about the root path that the indexing was started from. The same paths from nested roots must
        # not be mixed up.
        tmpl_roots = ["/", "/aX", "/aX/bY", "/aX/bY/cZ"]
        root_prefixes = [
            "",  # posix
            "C:",  # nt - normal drive letter
            "//localHost/c$",  # nt - UNC path
        ]
        rel_paths = ["0", "1/2", "1/3/4", "5/6", "5/7", "5/8/9"]
        # check self consistency for path generation
        for i in range(1, len(tmpl_roots)):
            for j in range(0, i):
                assert tmpl_roots[i].startswith(tmpl_roots[j])
        docs_r = dict()
        all_paths = set()
        children_map = defaultdict(dict)
        root_query_accessors = defaultdict(set)
        sub_query_accessors = defaultdict(set)
        for root_prefix, tmpl_root in itertools.product(root_prefixes, tmpl_roots):
            root = f"{root_prefix}{tmpl_root}"
            for sub_root in tmpl_roots:
                if len(tmpl_root) > len(sub_root):
                    continue
                deepest_root = f"{root_prefix}{sub_root}"
                root_rel = deepest_root[len(root) :].lstrip("/")
                for rel_path in rel_paths:
                    if root_rel and rel_path:
                        rel = f"{root_rel}/{rel_path}"
                    elif root_rel:
                        rel = root_rel
                    elif rel_path:
                        rel = rel_path
                    else:
                        assert False, "relative path must be something for indexing (root is dir thus not indexed)"
                    final_path = f"{root}{'' if root.endswith('/') else '/'}{rel}"
                    if final_path not in docs_r:
                        docs_r[final_path] = {"path": final_path, "file": True, "accessors": set()}
                    docs_r[final_path]["accessors"].add(create_accessor_path(root, rel))
                    all_paths.add(final_path)
                    parent_path, _, doc_name = create_accessor_path(root, rel).rpartition("/")
                    assert create_accessor_path(root, rel.rpartition("/")[0]) == parent_path
                    assert rel.rpartition("/")[2] == doc_name
                    children_map[parent_path][doc_name] = final_path
                    qr = create_accessor_path(root)
                    root_query_accessors[qr].add(final_path)
                    assert qr == create_accessor_path(root, "")
                    rel_parts = rel.split("/")
                    for i in range(1, len(rel_parts)):
                        sub_rel = "/".join(rel_parts[:i])
                        sub_query_accessors[create_accessor_path(root, sub_rel)].add(final_path)
        # NOTE: this is needed because of an implementation deficiency; see NOTE in create_accessor_path()
        root_query_accessors["/"].update(root_query_accessors["//localHost/c$/"])
        # insert docs with nested roots
        docs_r = list(docs_r.values())
        for d in docs_r:
            d["accessors"] = sorted(d["accessors"])
        await run_check_update(session, u_url, auth_h, docs_r)
        # test tree queries
        for q_set in [root_query_accessors, sub_query_accessors]:
            for root_acc, paths in q_set.items():
                q_paths = set(
                    await run_check_query(
                        session,
                        q_url,
                        auth_h,
                        [
                            ("q", create_docs_query_tree(root_acc)),
                            ("fl", "path"),
                            ("start", 0),
                            ("rows", len(all_paths) + 10),
                        ],
                        return_paths=True,
                    )
                )
                assert paths == q_paths
        # test children and exact queries
        for parent_path, name_paths in children_map.items():
            q_paths = set(
                await run_check_query(
                    session,
                    q_url,
                    auth_h,
                    [
                        ("q", create_docs_query_children(parent_path)),
                        ("fl", "path"),
                        ("start", 0),
                        ("rows", len(all_paths) + 10),
                    ],
                    return_paths=True,
                )
            )
            assert set(p for p in name_paths.values()) == q_paths, f"{parent_path}\n{name_paths}\n{q_paths}"
            for name, path in name_paths.items():
                q_paths = set(
                    await run_check_query(
                        session,
                        q_url,
                        auth_h,
                        [
                            ("q", create_docs_query_exact(parent_path, name)),
                            ("fl", "path"),
                            ("start", 0),
                            ("rows", len(all_paths) + 10),
                        ],
                        return_paths=True,
                    )
                )
                assert {path} == q_paths

        # check indexing and querying root documents (some are edge cases)
        fs_root_paths = (
            ["//?/c:/", "//localhost/c$/", "//?/UNC/localhost/c$/"] if sys.platform.startswith("win") else ["/"]
        )
        docs2_roots = sorted(
            itertools.chain(
                *[[fs_root_path, f"{fs_root_path}a", f"{fs_root_path}a/b"] for fs_root_path in fs_root_paths]
            )
        )

        def docs2_query_args(*roots_):
            return [
                ("q", " OR ".join([create_docs_query_tree(create_accessor_path(r_)) for r_ in roots_])),
                ("fl", "path,file,accessors,mtime"),
                ("start", 0),
                ("rows", 2000),
                ("sort", "path asc"),
            ]

        # clear indexed docs for next tests
        prev_paths = await run_check_query(session, q_url, auth_h, docs2_query_args(*docs2_roots), return_paths=True)
        await run_check_update(
            session, u_url, auth_h, serialize_solr_update_json([("delete", {"id": pp}) for pp in prev_paths])
        )
        assert await run_check_query(session, q_url, auth_h, docs2_query_args(*docs2_roots), return_paths=True) == []
        docs2 = {
            docs2_root: [
                {
                    "path": docs2_root,
                    "file": False,
                    "accessors": [create_accessor_path(docs2_root)],
                    "mtime": 1649757718358567545 + i * 2,
                },
                {
                    "path": f"{docs2_root}/x",
                    "file": False,
                    "accessors": [create_accessor_path(docs2_root, "x")],
                    "mtime": 1649757718358567545 + i * 2 + 1,
                },
            ]
            for i, docs2_root in enumerate(docs2_roots)
        }
        docs2_flat = sorted(itertools.chain(*list(docs2.values())), key=lambda t_: t_["path"])
        await run_check_update(session, u_url, auth_h, docs2_flat)
        for root, docs in docs2.items():
            await run_check_query(session, q_url, auth_h, docs2_query_args(root), docs, len(docs))
        await run_check_query(session, q_url, auth_h, docs2_query_args(*docs2_roots), docs2_flat, len(docs2_flat))
        # exact queries on root has the most relevant test for the edge-cases
        for root, docs in docs2.items():
            for doc in docs:
                name = "" if doc["path"] == root else doc["path"][len(root) + 1 :]
                await run_check_query(
                    session,
                    q_url,
                    auth_h,
                    [
                        ("q", create_docs_query_exact(create_accessor_path(root), name)),
                        ("fl", "path,file,accessors,mtime"),
                        ("start", 0),
                        ("rows", 5),
                    ],
                    [doc],
                    1,
                )

        # check scope-behaviour of commitWithin & soft commit interval
        docs3 = [
            {
                "path": "/var/xy/xy",
                "file": True,
                "accessors": [create_accessor_path("/var/xy", "xy")],
                "size": 6200,
                "mtime": 4444,
            },
            {
                "path": "/var/xy/az",
                "file": True,
                "accessors": [create_accessor_path("/var/xy", "az")],
                "size": 6201,
                "mtime": 4445,
            },
            {
                "path": "/var/xy/gh",
                "file": True,
                "accessors": [create_accessor_path("/var/xy", "gh")],
                "size": 6202,
                "mtime": 4446,
            },
            {
                "path": "/var/xy/qw",
                "file": True,
                "accessors": [create_accessor_path("/var/xy", "qw")],
                "size": 6203,
                "mtime": 4447,
            },
            {
                "path": "/var/xy/zh",
                "file": True,
                "accessors": [create_accessor_path("/var/xy", "zh")],
                "size": 6204,
                "mtime": 4448,
            },
            {
                "path": "/var/xy/zx",
                "file": True,
                "accessors": [create_accessor_path("/var/xy", "zx")],
                "size": 6205,
                "mtime": 4449,
            },
        ]

        async def query_docs3():
            return set(
                await run_check_query(
                    session,
                    q_url,
                    auth_h,
                    [
                        ("q", create_docs_query_children(create_accessor_path("/var/xy"))),
                        ("fl", "path"),
                        ("start", 0),
                        ("rows", 10),
                    ],
                    return_paths=True,
                )
            )

        # cannot reset or coordinate with commit timers on solr's end, so if something is amiss try again once
        for attempt in range(2):
            try:
                await run_check_update(session, u_url, auth_h, {"add": {"doc": docs3[3]}}, commit=False)
                await run_check_update(session, u_url, auth_h, {}, commit=False, soft_commit=True)
                assert {docs3[3]["path"]} == await query_docs3()
                await run_check_update(
                    session, u_url, auth_h, {"add": {"doc": docs3[0]}, "delete": {"id": docs3[3]["path"]}}, commit=False
                )
                assert {docs3[3]["path"]} == await query_docs3()
                await run_check_update(
                    session, u_url, auth_h, {"add": {"doc": docs3[1]}}, commit=False, commit_within=2000
                )
                assert {docs3[3]["path"]} == await query_docs3()
                await asyncio.sleep(0.25)
                assert {docs3[3]["path"]} == await query_docs3()
                await run_check_update(session, u_url, auth_h, [docs3[2]], commit=False)
                await run_check_update(
                    session,
                    u_url,
                    auth_h,
                    serialize_solr_update_json([("add", docs3[4]), ("add", docs3[5])]),
                    commit=False,
                )
                assert {docs3[3]["path"]} == await query_docs3()
                await asyncio.sleep(0.25)
                assert {docs3[3]["path"]} == await query_docs3()
                await asyncio.sleep(1.75)
                assert {docs3[i]["path"] for i in range(len(docs3)) if i != 3} == await query_docs3()
                # see if soft-commit interval is correctly ~15 seconds
                soft_commit_interval = 15.0
                await run_check_update(session, u_url, auth_h, {"add": {"doc": docs3[3]}}, commit=False)
                await asyncio.sleep(soft_commit_interval - 1.0)
                assert {docs3[i]["path"] for i in range(len(docs3)) if i != 3} == await query_docs3()
                await asyncio.sleep(1.5)
                assert {docs3[i]["path"] for i in range(len(docs3))} == await query_docs3()
                break
            except AssertionError:
                if attempt != 0:
                    raise
                await run_check_update(
                    session,
                    u_url,
                    auth_h,
                    serialize_solr_update_json([("delete", {"id": d["path"]}) for d in docs3]),
                    commit=False,
                    soft_commit=True,
                )
                assert set() == await query_docs3()
