from __future__ import annotations

import argparse
import json
import re
import unicodedata
from typing import Any, Iterable


GRAPH_SCOPE = "zpwiki"
DEFAULT_LIMIT = 40
DEFAULT_MAX_ENTITIES = 8

RELATIONSHIP_TYPES = (
    "HAS_DOCUMENT|AUTHORED|IN_CATEGORY|HAS_KEYWORD|"
    "DESCRIBES|HAS_WORK|EVIDENCED_BY|AUTHORED_DOCUMENT"
)

CATALOG_QUERY = """
MATCH (n)
WHERE n.graph_scope = $scope
  AND (
    n:Person
    OR n:Thesis
    OR n:Work
    OR n:Keyword
    OR n:Topic
    OR n:Category
    OR n:DocumentAuthor
  )
RETURN
  elementId(n) AS element_id,
  CASE
    WHEN n:Person THEN 'Person'
    WHEN n:Thesis THEN 'Thesis'
    WHEN n:Work THEN 'Work'
    WHEN n:Keyword THEN 'Keyword'
    WHEN n:Topic THEN 'Topic'
    WHEN n:Category THEN 'Category'
    WHEN n:DocumentAuthor THEN 'DocumentAuthor'
    ELSE 'Other'
  END AS kind,
  n.id AS id,
  n.name AS name,
  n.title AS title,
  n.heading AS heading,
  n.start_year AS start_year,
  n.year AS year,
  n.thesis_type AS thesis_type,
  n.thesis_type_label AS thesis_type_label,
  n.work_type AS work_type,
  n.work_type_label AS work_type_label,
  n.person_kind AS person_kind,
  n.source_document AS source_document
ORDER BY kind, name, title, heading
"""

NEIGHBORHOOD_QUERY = f"""
MATCH (start)
WHERE elementId(start) IN $start_ids
MATCH p=(start)-[:{RELATIONSHIP_TYPES}*1..2]-(target)
WHERE all(node IN nodes(p) WHERE node.graph_scope = $scope)
  AND (
    size($target_kinds) = 0
    OR any(
      node IN nodes(p)
      WHERE
        ('Person' IN $target_kinds AND node:Person)
        OR ('Thesis' IN $target_kinds AND node:Thesis)
        OR ('Work' IN $target_kinds AND node:Work)
        OR ('Keyword' IN $target_kinds AND node:Keyword)
        OR ('Topic' IN $target_kinds AND node:Topic)
        OR ('Category' IN $target_kinds AND node:Category)
        OR (
          'DocumentAuthor' IN $target_kinds
          AND node:DocumentAuthor
        )
        OR ('Document' IN $target_kinds AND node:Document)
    )
  )
RETURN
  [node IN nodes(p) | {{
    element_id: elementId(node),
    kind: CASE
      WHEN node:Person THEN 'Person'
      WHEN node:Thesis THEN 'Thesis'
      WHEN node:Work THEN 'Work'
      WHEN node:Document THEN 'Document'
      WHEN node:Keyword THEN 'Keyword'
      WHEN node:Topic THEN 'Topic'
      WHEN node:Category THEN 'Category'
      WHEN node:DocumentAuthor THEN 'DocumentAuthor'
      ELSE 'Other'
    END,
    id: node.id,
    name: node.name,
    title: node.title,
    heading: node.heading,
    path: node.path,
    start_year: node.start_year,
    year: node.year,
    thesis_type: node.thesis_type,
    thesis_type_label: node.thesis_type_label,
    work_type: node.work_type,
    work_type_label: node.work_type_label,
    person_kind: node.person_kind,
    source_document: node.source_document,
    source_url: node.source_url
  }}] AS nodes,
  [relationship IN relationships(p) |
    type(relationship)
  ] AS relationships,
  length(p) AS hops
ORDER BY hops
LIMIT $limit
"""

PAIR_PATH_QUERY = f"""
MATCH (left), (right)
WHERE elementId(left) IN $start_ids
  AND elementId(right) IN $start_ids
  AND elementId(left) < elementId(right)
MATCH p=allShortestPaths(
  (left)-[:{RELATIONSHIP_TYPES}*1..4]-(right)
)
WHERE all(node IN nodes(p) WHERE node.graph_scope = $scope)
RETURN
  [node IN nodes(p) | {{
    element_id: elementId(node),
    kind: CASE
      WHEN node:Person THEN 'Person'
      WHEN node:Thesis THEN 'Thesis'
      WHEN node:Work THEN 'Work'
      WHEN node:Document THEN 'Document'
      WHEN node:Keyword THEN 'Keyword'
      WHEN node:Topic THEN 'Topic'
      WHEN node:Category THEN 'Category'
      WHEN node:DocumentAuthor THEN 'DocumentAuthor'
      ELSE 'Other'
    END,
    id: node.id,
    name: node.name,
    title: node.title,
    heading: node.heading,
    path: node.path,
    start_year: node.start_year,
    year: node.year,
    thesis_type: node.thesis_type,
    thesis_type_label: node.thesis_type_label,
    work_type: node.work_type,
    work_type_label: node.work_type_label,
    person_kind: node.person_kind,
    source_document: node.source_document,
    source_url: node.source_url
  }}] AS nodes,
  [relationship IN relationships(p) |
    type(relationship)
  ] AS relationships,
  length(p) AS hops
ORDER BY hops
LIMIT $limit
"""

THESIS_FACTS_QUERY = """
MATCH
  (person:Person)
  -[:AUTHORED]->
  (thesis:Thesis)
  -[:EVIDENCED_BY]->
  (document:Document)
WHERE person.graph_scope = $scope
  AND thesis.graph_scope = $scope
  AND document.graph_scope = $scope
  AND (
    size($person_ids) = 0
    OR elementId(person) IN $person_ids
  )
  AND (
    size($thesis_ids) = 0
    OR elementId(thesis) IN $thesis_ids
  )
  AND (
    size($years) = 0
    OR thesis.year IN $years
  )
  AND (
    size($thesis_types) = 0
    OR thesis.thesis_type IN $thesis_types
  )
RETURN
  elementId(person) AS person_element_id,
  person.name AS person,
  person.start_year AS start_year,
  elementId(thesis) AS thesis_element_id,
  thesis.id AS thesis_id,
  thesis.title AS title,
  thesis.heading AS heading,
  thesis.year AS year,
  thesis.thesis_type AS thesis_type,
  thesis.thesis_type_label AS thesis_type_label,
  document.path AS document_path,
  document.source_url AS source_url,
  thesis.source_chunk_ids AS source_chunk_ids
ORDER BY person, year, thesis_type, title
LIMIT $limit
"""

THESIS_PATHS_QUERY = """
MATCH p=
  (person:Person)
  -[:AUTHORED]->
  (thesis:Thesis)
  -[:EVIDENCED_BY]->
  (document:Document)
WHERE person.graph_scope = $scope
  AND thesis.graph_scope = $scope
  AND document.graph_scope = $scope
  AND (
    size($person_ids) = 0
    OR elementId(person) IN $person_ids
  )
  AND (
    size($thesis_ids) = 0
    OR elementId(thesis) IN $thesis_ids
  )
  AND (
    size($years) = 0
    OR thesis.year IN $years
  )
  AND (
    size($thesis_types) = 0
    OR thesis.thesis_type IN $thesis_types
  )
RETURN
  [node IN nodes(p) | {
    element_id: elementId(node),
    kind: CASE
      WHEN node:Person THEN 'Person'
      WHEN node:Thesis THEN 'Thesis'
      WHEN node:Document THEN 'Document'
      ELSE 'Other'
    END,
    id: node.id,
    name: node.name,
    title: node.title,
    heading: node.heading,
    path: node.path,
    start_year: node.start_year,
    year: node.year,
    thesis_type: node.thesis_type,
    thesis_type_label: node.thesis_type_label,
    source_document: node.source_document,
    source_url: node.source_url
  }] AS nodes,
  [relationship IN relationships(p) |
    type(relationship)
  ] AS relationships,
  length(p) AS hops
ORDER BY
  person.name,
  thesis.year,
  thesis.thesis_type,
  thesis.title
LIMIT $limit
"""

WORK_FACTS_QUERY = """
MATCH
  (person:Person)
  -[:HAS_WORK]->
  (work:Work)
  -[:EVIDENCED_BY]->
  (document:Document)
WHERE person.graph_scope = $scope
  AND work.graph_scope = $scope
  AND document.graph_scope = $scope
  AND (
    size($person_ids) = 0
    OR elementId(person) IN $person_ids
  )
  AND (
    size($work_ids) = 0
    OR elementId(work) IN $work_ids
  )
  AND (
    size($years) = 0
    OR work.year IN $years
  )
  AND (
    size($work_types) = 0
    OR work.work_type IN $work_types
  )
RETURN
  elementId(person) AS person_element_id,
  person.name AS person,
  person.start_year AS start_year,
  elementId(work) AS work_element_id,
  work.id AS work_id,
  work.title AS title,
  work.heading AS heading,
  work.year AS year,
  work.work_type AS work_type,
  work.work_type_label AS work_type_label,
  document.path AS document_path,
  document.source_url AS source_url,
  work.source_chunk_ids AS source_chunk_ids
ORDER BY person, year, work_type, title
LIMIT $limit
"""

WORK_PATHS_QUERY = """
MATCH p=
  (person:Person)
  -[:HAS_WORK]->
  (work:Work)
  -[:EVIDENCED_BY]->
  (document:Document)
WHERE person.graph_scope = $scope
  AND work.graph_scope = $scope
  AND document.graph_scope = $scope
  AND (
    size($person_ids) = 0
    OR elementId(person) IN $person_ids
  )
  AND (
    size($work_ids) = 0
    OR elementId(work) IN $work_ids
  )
  AND (
    size($years) = 0
    OR work.year IN $years
  )
  AND (
    size($work_types) = 0
    OR work.work_type IN $work_types
  )
RETURN
  [node IN nodes(p) | {
    element_id: elementId(node),
    kind: CASE
      WHEN node:Person THEN 'Person'
      WHEN node:Work THEN 'Work'
      WHEN node:Document THEN 'Document'
      ELSE 'Other'
    END,
    id: node.id,
    name: node.name,
    title: node.title,
    heading: node.heading,
    path: node.path,
    start_year: node.start_year,
    year: node.year,
    work_type: node.work_type,
    work_type_label: node.work_type_label,
    source_document: node.source_document,
    source_url: node.source_url
  }] AS nodes,
  [relationship IN relationships(p) |
    type(relationship)
  ] AS relationships,
  length(p) AS hops
ORDER BY
  person.name,
  work.year,
  work.work_type,
  work.title
LIMIT $limit
"""

START_YEAR_FACTS_QUERY = """
MATCH (person:Person)
WHERE person.graph_scope = $scope
  AND (
    size($person_ids) = 0
    OR elementId(person) IN $person_ids
  )
  AND (
    size($years) = 0
    OR person.start_year IN $years
  )
OPTIONAL MATCH (document:Document)
WHERE document.graph_scope = $scope
  AND document.path = person.source_document
RETURN
  elementId(person) AS person_element_id,
  person.name AS person,
  person.start_year AS start_year,
  person.source_document AS document_path,
  document.source_url AS source_url
ORDER BY person
LIMIT $limit
"""

PERSON_KEYWORD_FACTS_QUERY = """
MATCH
  (person:Person)
  -[:HAS_DOCUMENT]->
  (document:Document)
  -[:HAS_KEYWORD]->
  (keyword:Keyword)
WHERE person.graph_scope = $scope
  AND document.graph_scope = $scope
  AND keyword.graph_scope = $scope
  AND (
    size($person_ids) = 0
    OR elementId(person) IN $person_ids
  )
  AND (
    size($keyword_ids) = 0
    OR elementId(keyword) IN $keyword_ids
  )
WITH
  person,
  keyword,
  collect(DISTINCT document) AS documents
WITH
  person,
  keyword,
  head(
    [
      document IN documents
      WHERE document.path = person.source_document
    ]
    + documents
  ) AS document
RETURN
  elementId(person) AS person_element_id,
  person.name AS person,
  elementId(keyword) AS keyword_element_id,
  keyword.id AS keyword_id,
  keyword.name AS keyword,
  document.path AS document_path,
  document.source_url AS source_url
ORDER BY person, keyword
LIMIT $limit
"""

PERSON_CATEGORY_FACTS_QUERY = """
MATCH
  (person:Person)
  -[:HAS_DOCUMENT]->
  (document:Document)
  -[:IN_CATEGORY]->
  (category:Category)
WHERE person.graph_scope = $scope
  AND document.graph_scope = $scope
  AND category.graph_scope = $scope
  AND (
    size($person_ids) = 0
    OR elementId(person) IN $person_ids
  )
  AND (
    size($category_ids) = 0
    OR elementId(category) IN $category_ids
  )
WITH
  person,
  category,
  collect(DISTINCT document) AS documents
WITH
  person,
  category,
  head(
    [
      document IN documents
      WHERE document.path = person.source_document
    ]
    + documents
  ) AS document
RETURN
  elementId(person) AS person_element_id,
  person.name AS person,
  elementId(category) AS category_element_id,
  category.name AS category,
  category.category_type AS category_type,
  category.year AS category_year,
  document.path AS document_path,
  document.source_url AS source_url
ORDER BY person, category
LIMIT $limit
"""

SHARED_KEYWORD_FACTS_QUERY = """
MATCH
  (left:Person)
  -[:HAS_DOCUMENT]->
  (left_document:Document)
  -[:HAS_KEYWORD]->
  (keyword:Keyword)
MATCH
  (right:Person)
  -[:HAS_DOCUMENT]->
  (right_document:Document)
  -[:HAS_KEYWORD]->
  (keyword)
WHERE left.graph_scope = $scope
  AND right.graph_scope = $scope
  AND left_document.graph_scope = $scope
  AND right_document.graph_scope = $scope
  AND keyword.graph_scope = $scope
  AND elementId(left) IN $person_ids
  AND elementId(right) IN $person_ids
  AND elementId(left) < elementId(right)
WITH
  left,
  right,
  keyword,
  collect(DISTINCT left_document) AS left_documents,
  collect(DISTINCT right_document) AS right_documents
WITH
  left,
  right,
  keyword,
  head(
    [
      document IN left_documents
      WHERE document.path = left.source_document
    ]
    + left_documents
  ) AS left_document,
  head(
    [
      document IN right_documents
      WHERE document.path = right.source_document
    ]
    + right_documents
  ) AS right_document
RETURN
  elementId(left) AS left_person_element_id,
  left.name AS left_person,
  left_document.path AS left_document_path,
  left_document.source_url AS left_source_url,
  elementId(right) AS right_person_element_id,
  right.name AS right_person,
  right_document.path AS right_document_path,
  right_document.source_url AS right_source_url,
  elementId(keyword) AS keyword_element_id,
  keyword.id AS keyword_id,
  keyword.name AS keyword
ORDER BY left_person, right_person, keyword
LIMIT $limit
"""

SHARED_CATEGORY_FACTS_QUERY = """
MATCH
  (left:Person)
  -[:HAS_DOCUMENT]->
  (left_document:Document)
  -[:IN_CATEGORY]->
  (category:Category)
MATCH
  (right:Person)
  -[:HAS_DOCUMENT]->
  (right_document:Document)
  -[:IN_CATEGORY]->
  (category)
WHERE left.graph_scope = $scope
  AND right.graph_scope = $scope
  AND left_document.graph_scope = $scope
  AND right_document.graph_scope = $scope
  AND category.graph_scope = $scope
  AND elementId(left) IN $person_ids
  AND elementId(right) IN $person_ids
  AND elementId(left) < elementId(right)
WITH
  left,
  right,
  category,
  collect(DISTINCT left_document) AS left_documents,
  collect(DISTINCT right_document) AS right_documents
WITH
  left,
  right,
  category,
  head(
    [
      document IN left_documents
      WHERE document.path = left.source_document
    ]
    + left_documents
  ) AS left_document,
  head(
    [
      document IN right_documents
      WHERE document.path = right.source_document
    ]
    + right_documents
  ) AS right_document
RETURN
  elementId(left) AS left_person_element_id,
  left.name AS left_person,
  left_document.path AS left_document_path,
  left_document.source_url AS left_source_url,
  elementId(right) AS right_person_element_id,
  right.name AS right_person,
  right_document.path AS right_document_path,
  right_document.source_url AS right_source_url,
  elementId(category) AS category_element_id,
  category.name AS category,
  category.category_type AS category_type,
  category.year AS category_year
ORDER BY left_person, right_person, category
LIMIT $limit
"""

SHARED_DOCUMENT_AUTHOR_FACTS_QUERY = """
MATCH
  (author:DocumentAuthor)
  -[:AUTHORED_DOCUMENT]->
  (left_document:Document)
  <-[:HAS_DOCUMENT]-
  (left:Person)
MATCH
  (author)
  -[:AUTHORED_DOCUMENT]->
  (right_document:Document)
  <-[:HAS_DOCUMENT]-
  (right:Person)
WHERE author.graph_scope = $scope
  AND left.graph_scope = $scope
  AND right.graph_scope = $scope
  AND left_document.graph_scope = $scope
  AND right_document.graph_scope = $scope
  AND elementId(left) IN $person_ids
  AND elementId(right) IN $person_ids
  AND elementId(left) < elementId(right)
WITH
  left,
  right,
  author,
  collect(DISTINCT left_document) AS left_documents,
  collect(DISTINCT right_document) AS right_documents
WITH
  left,
  right,
  author,
  head(
    [
      document IN left_documents
      WHERE document.path = left.source_document
    ]
    + left_documents
  ) AS left_document,
  head(
    [
      document IN right_documents
      WHERE document.path = right.source_document
    ]
    + right_documents
  ) AS right_document
RETURN
  elementId(left) AS left_person_element_id,
  left.name AS left_person,
  left_document.path AS left_document_path,
  left_document.source_url AS left_source_url,
  elementId(right) AS right_person_element_id,
  right.name AS right_person,
  right_document.path AS right_document_path,
  right_document.source_url AS right_source_url,
  elementId(author) AS document_author_element_id,
  author.id AS document_author_id,
  author.name AS document_author
ORDER BY left_person, right_person, document_author
LIMIT $limit
"""


def normalize_for_match(
    value: str,
) -> str:
    decomposed = unicodedata.normalize(
        "NFKD",
        str(value),
    )

    without_marks = "".join(
        character
        for character in decomposed
        if not unicodedata.combining(
            character
        )
    )

    normalized = re.sub(
        r"[^a-zA-Z0-9]+",
        " ",
        without_marks.casefold(),
    )

    return re.sub(
        r"\s+",
        " ",
        normalized,
    ).strip()


def normalized_tokens(
    value: str,
) -> list[str]:
    normalized = normalize_for_match(
        value
    )

    if not normalized:
        return []

    return normalized.split()


def has_stem(
    tokens: Iterable[str],
    stem: str,
) -> bool:
    return any(
        token.startswith(
            stem
        )
        for token in tokens
    )


def extract_query_constraints(
    query: str,
) -> dict[str, Any]:
    normalized = normalize_for_match(
        query
    )

    tokens = normalized.split()
    token_set = set(
        tokens
    )

    years = sorted(
        {
            int(value)
            for value in re.findall(
                r"\b(20\d{2})\b",
                normalized,
            )
        }
    )

    has_work_word = has_stem(
        tokens,
        "prac",
    )

    has_project_word = has_stem(
        tokens,
        "projekt",
    )

    thesis_types: list[str] = []
    work_types: list[str] = []

    if (
        (
            has_stem(
                tokens,
                "bakalar",
            )
            and has_work_word
        )
        or "bp" in token_set
    ):
        thesis_types.append(
            "bachelor_thesis"
        )

    if (
        (
            has_stem(
                tokens,
                "diplom",
            )
            and has_work_word
        )
        or "dp" in token_set
    ):
        thesis_types.append(
            "diploma_thesis"
        )

    if (
        (
            has_stem(
                tokens,
                "dizert",
            )
            and has_work_word
        )
        or "phd" in token_set
    ):
        thesis_types.append(
            "doctoral_thesis"
        )

    if (
        has_stem(
            tokens,
            "diplom",
        )
        and has_project_word
    ):
        work_types.append(
            "diploma_project"
        )

    if (
        has_stem(
            tokens,
            "timov",
        )
        and has_project_word
    ):
        work_types.append(
            "team_project"
        )

    if (
        has_stem(
            tokens,
            "vyskum",
        )
        and has_project_word
    ):
        work_types.append(
            "research_project"
        )

    if "tp" in token_set:
        work_types.append(
            "team_project"
        )

    if "vp" in token_set:
        work_types.append(
            "research_project"
        )

    thesis_types = list(
        dict.fromkeys(
            thesis_types
        )
    )

    work_types = list(
        dict.fromkeys(
            work_types
        )
    )

    mentions_start_year = (
        has_stem(
            tokens,
            "zaciat",
        )
        and has_stem(
            tokens,
            "stud",
        )
    )

    mentions_keyword = (
        has_stem(
            tokens,
            "kluc",
        )
        or "tag" in token_set
        or "tagy" in token_set
    )

    mentions_topic = (
        has_stem(
            tokens,
            "tem",
        )
    )

    mentions_category = (
        has_stem(
            tokens,
            "kategori",
        )
    )

    mentions_document_author = (
        has_stem(
            tokens,
            "autor",
        )
        and (
            has_stem(
                tokens,
                "dokument",
            )
            or has_stem(
                tokens,
                "strank",
            )
            or "wiki" in token_set
        )
    )

    mentions_title = (
        has_stem(
            tokens,
            "nazov",
        )
    )

    mentions_shared = (
        has_stem(
            tokens,
            "spoloc",
        )
        or has_stem(
            tokens,
            "rovnak",
        )
    )

    mentions_person_answer = (
        "kto" in token_set
        or (
            has_stem(
                tokens,
                "ktor",
            )
            and (
                has_stem(
                    tokens,
                    "student",
                )
                or has_stem(
                    tokens,
                    "osob",
                )
                or has_stem(
                    tokens,
                    "clov",
                )
            )
        )
    )

    mentions_thesis = (
        bool(
            thesis_types
        )
        or (
            has_stem(
                tokens,
                "zaverec",
            )
            and has_work_word
        )
    )

    mentions_work = bool(
        work_types
    )

    target_kinds: list[str] = []

    if mentions_keyword:
        target_kinds.append(
            "Keyword"
        )

    if mentions_topic:
        target_kinds.extend(
            [
                "Keyword",
                "Topic",
            ]
        )

    if mentions_category:
        target_kinds.append(
            "Category"
        )

    if mentions_document_author:
        target_kinds.append(
            "DocumentAuthor"
        )

    if mentions_thesis:
        target_kinds.append(
            "Thesis"
        )

    if mentions_work:
        target_kinds.append(
            "Work"
        )

    if (
        mentions_title
        and not mentions_thesis
        and not mentions_work
    ):
        target_kinds.extend(
            [
                "Thesis",
                "Work",
            ]
        )

    if (
        years
        and not mentions_start_year
        and not mentions_thesis
        and not mentions_work
    ):
        target_kinds.extend(
            [
                "Thesis",
                "Work",
            ]
        )

    if (
        mentions_person_answer
        and not mentions_shared
    ):
        target_kinds = [
            "Person"
        ]

    return {
        "years": years,
        "thesis_types": (
            thesis_types
        ),
        "work_types": (
            work_types
        ),
        "mentions_start_year": (
            mentions_start_year
        ),
        "mentions_keyword": (
            mentions_keyword
        ),
        "mentions_topic": (
            mentions_topic
        ),
        "mentions_category": (
            mentions_category
        ),
        "mentions_document_author": (
            mentions_document_author
        ),
        "mentions_title": (
            mentions_title
        ),
        "mentions_shared": (
            mentions_shared
        ),
        "mentions_person_answer": (
            mentions_person_answer
        ),
        "mentions_thesis": (
            mentions_thesis
        ),
        "mentions_work": (
            mentions_work
        ),
        "target_kinds": list(
            dict.fromkeys(
                target_kinds
            )
        ),
    }


def entity_aliases(
    entity: dict[str, Any],
) -> list[str]:
    kind = str(
        entity.get(
            "kind"
        )
        or ""
    )

    values: list[Any] = []

    if kind in {
        "Person",
        "DocumentAuthor",
    }:
        values.append(
            entity.get(
                "name"
            )
        )

    elif kind in {
        "Keyword",
        "Topic",
    }:
        values.extend(
            [
                entity.get(
                    "name"
                ),
                entity.get(
                    "id"
                ),
            ]
        )

    elif kind == "Category":
        values.append(
            entity.get(
                "name"
            )
        )

    elif kind in {
        "Thesis",
        "Work",
    }:
        values.append(
            entity.get(
                "title"
            )
        )

    aliases: list[str] = []

    for value in values:
        if value is None:
            continue

        normalized = normalize_for_match(
            str(value)
        )

        if normalized:
            aliases.append(
                normalized
            )

    return list(
        dict.fromkeys(
            aliases
        )
    )


def tokens_match(
    left: str,
    right: str,
) -> bool:
    if left == right:
        return True

    shorter = min(
        len(left),
        len(right),
    )

    if shorter <= 2:
        return False

    required = (
        3
        if shorter <= 4
        else 4
        if shorter <= 6
        else 5
    )

    return (
        left[:required]
        == right[:required]
    )


def alias_match_score(
    normalized_query: str,
    query_tokens: set[str],
    alias: str,
) -> float:
    alias_tokens = (
        alias.split()
    )

    if not alias_tokens:
        return 0.0

    if len(
        alias_tokens
    ) == 1:
        if alias not in query_tokens:
            return 0.0

        return (
            50.0
            + min(
                len(alias),
                20,
            )
            / 100.0
        )

    padded_query = (
        f" {normalized_query} "
    )

    padded_alias = (
        f" {alias} "
    )

    if (
        padded_alias
        in padded_query
    ):
        score = (
            100.0
            + len(
                alias_tokens
            )
            * 5.0
            + min(
                len(alias),
                100,
            )
            / 100.0
        )

        if (
            normalized_query
            == alias
        ):
            score += 50.0

        return score

    if all(
        any(
            tokens_match(
                alias_token,
                query_token,
            )
            for query_token
            in query_tokens
        )
        for alias_token
        in alias_tokens
    ):
        return (
            95.0
            + len(
                alias_tokens
            )
            * 5.0
            + min(
                len(alias),
                100,
            )
            / 100.0
        )

    return 0.0


def suppress_nested_entity_matches(
    matches: list[
        dict[str, Any]
    ],
) -> list[dict[str, Any]]:
    result: list[
        dict[str, Any]
    ] = []

    for item in matches:
        kind = str(
            item.get(
                "kind"
            )
            or ""
        )

        if kind not in {
            "Thesis",
            "Work",
        }:
            result.append(
                item
            )
            continue

        alias = str(
            item.get(
                "matched_alias"
            )
            or ""
        )

        score = float(
            item.get(
                "match_score"
            )
            or 0.0
        )

        dominated = False

        for other in matches:
            if other is item:
                continue

            if (
                other.get(
                    "kind"
                )
                != kind
            ):
                continue

            other_alias = str(
                other.get(
                    "matched_alias"
                )
                or ""
            )

            other_score = float(
                other.get(
                    "match_score"
                )
                or 0.0
            )

            if (
                not alias
                or not other_alias
                or alias == other_alias
            ):
                continue

            if (
                alias in other_alias
                and len(
                    other_alias
                )
                > len(
                    alias
                )
                and other_score
                > score
            ):
                dominated = True
                break

        if not dominated:
            result.append(
                item
            )

    return result


def resolve_entities(
    query: str,
    catalog: list[
        dict[str, Any]
    ],
    *,
    max_entities: int = (
        DEFAULT_MAX_ENTITIES
    ),
    preferred_kinds: Iterable[
        str
    ] = (),
) -> list[dict[str, Any]]:
    normalized_query = (
        normalize_for_match(
            query
        )
    )

    query_tokens = set(
        normalized_query.split()
    )

    preferred = set(
        preferred_kinds
    )

    matches: list[
        dict[str, Any]
    ] = []

    for entity in catalog:
        best_score = 0.0
        best_alias = ""

        for alias in entity_aliases(
            entity
        ):
            score = (
                alias_match_score(
                    normalized_query,
                    query_tokens,
                    alias,
                )
            )

            if (
                score > 0
                and entity.get(
                    "kind"
                )
                in preferred
            ):
                score += 2.0

            if score > best_score:
                best_score = score
                best_alias = alias

        if best_score <= 0.0:
            continue

        item = dict(
            entity
        )

        item[
            "matched_alias"
        ] = best_alias

        item[
            "match_score"
        ] = round(
            best_score,
            3,
        )

        matches.append(
            item
        )

    person_aliases = {
        str(
            item.get(
                "matched_alias"
            )
            or ""
        )
        for item in matches
        if item.get(
            "kind"
        )
        == "Person"
    }

    matches = [
        item
        for item in matches
        if not (
            item.get(
                "kind"
            )
            == "DocumentAuthor"
            and str(
                item.get(
                    "matched_alias"
                )
                or ""
            )
            in person_aliases
        )
    ]

    matches = (
        suppress_nested_entity_matches(
            matches
        )
    )

    matches.sort(
        key=lambda item: (
            -float(
                item[
                    "match_score"
                ]
            ),
            str(
                item.get(
                    "kind"
                )
                or ""
            ),
            str(
                item.get(
                    "name"
                )
                or item.get(
                    "title"
                )
                or item.get(
                    "id"
                )
                or ""
            ),
        )
    )

    selected: list[
        dict[str, Any]
    ] = []

    seen_ids: set[
        str
    ] = set()

    for item in matches:
        element_id = str(
            item.get(
                "element_id"
            )
            or ""
        )

        if not element_id:
            continue

        if element_id in seen_ids:
            continue

        seen_ids.add(
            element_id
        )

        selected.append(
            item
        )

        if (
            len(selected)
            >= max_entities
        ):
            break

    return selected


def rows_from_result(
    result: Any,
) -> list[dict[str, Any]]:
    rows: list[
        dict[str, Any]
    ] = []

    for record in result:
        if isinstance(
            record,
            dict,
        ):
            rows.append(
                dict(
                    record
                )
            )
            continue

        data = getattr(
            record,
            "data",
            None,
        )

        if callable(
            data
        ):
            rows.append(
                dict(
                    data()
                )
            )
            continue

        rows.append(
            dict(
                record
            )
        )

    return rows


def run_rows(
    session: Any,
    query: str,
    **parameters: Any,
) -> list[dict[str, Any]]:
    return rows_from_result(
        session.run(
            query,
            **parameters,
        )
    )


def deduplicate_paths(
    paths: list[
        dict[str, Any]
    ],
) -> list[dict[str, Any]]:
    result: list[
        dict[str, Any]
    ] = []

    seen: set[
        str
    ] = set()

    for path in paths:
        key = json.dumps(
            {
                "nodes": path.get(
                    "nodes",
                    [],
                ),
                "relationships": (
                    path.get(
                        "relationships",
                        [],
                    )
                ),
            },
            ensure_ascii=False,
            sort_keys=True,
            default=str,
        )

        if key in seen:
            continue

        seen.add(
            key
        )

        result.append(
            path
        )

    return result


def deduplicate_rows(
    rows: list[
        dict[str, Any]
    ],
    key_fields: tuple[
        str,
        ...,
    ],
) -> list[dict[str, Any]]:
    result: list[
        dict[str, Any]
    ] = []

    seen: set[
        tuple[str, ...]
    ] = set()

    for row in rows:
        key = tuple(
            str(
                row.get(
                    field
                )
                or ""
            )
            for field
            in key_fields
        )

        if key in seen:
            continue

        seen.add(
            key
        )

        result.append(
            row
        )

    return result


def collect_documents(
    resolved_entities: list[
        dict[str, Any]
    ],
    paths: list[
        dict[str, Any]
    ],
    facts: list[
        dict[str, Any]
    ],
) -> list[dict[str, Any]]:
    documents: dict[
        str,
        dict[str, Any],
    ] = {}

    def add_document(
        path: Any,
        source_url: Any = None,
    ) -> None:
        value = str(
            path
            or ""
        ).strip()

        if not value:
            return

        current = (
            documents.setdefault(
                value,
                {
                    "document_path": (
                        value
                    ),
                    "source_url": None,
                },
            )
        )

        if (
            source_url
            and not current.get(
                "source_url"
            )
        ):
            current[
                "source_url"
            ] = str(
                source_url
            )

    for entity in (
        resolved_entities
    ):
        add_document(
            entity.get(
                "source_document"
            )
        )

    for path in paths:
        for node in path.get(
            "nodes",
            [],
        ):
            if not isinstance(
                node,
                dict,
            ):
                continue

            if (
                node.get(
                    "kind"
                )
                == "Document"
            ):
                add_document(
                    node.get(
                        "path"
                    ),
                    node.get(
                        "source_url"
                    ),
                )

            add_document(
                node.get(
                    "source_document"
                )
            )

    for fact in facts:
        add_document(
            fact.get(
                "document_path"
            ),
            fact.get(
                "source_url"
            ),
        )

        add_document(
            fact.get(
                "left_document_path"
            ),
            fact.get(
                "left_source_url"
            ),
        )

        add_document(
            fact.get(
                "right_document_path"
            ),
            fact.get(
                "right_source_url"
            ),
        )

    return sorted(
        documents.values(),
        key=lambda item: (
            item[
                "document_path"
            ]
        ),
    )


def build_facts(
    thesis_rows: list[
        dict[str, Any]
    ],
    work_rows: list[
        dict[str, Any]
    ],
    start_year_rows: list[
        dict[str, Any]
    ],
    person_keyword_rows: list[
        dict[str, Any]
    ],
    person_category_rows: list[
        dict[str, Any]
    ],
    shared_keyword_rows: list[
        dict[str, Any]
    ],
    shared_category_rows: list[
        dict[str, Any]
    ],
    shared_document_author_rows: list[
        dict[str, Any]
    ],
    *,
    mentions_shared: bool,
) -> list[dict[str, Any]]:
    facts: list[
        dict[str, Any]
    ] = []

    for row in thesis_rows:
        fact = dict(
            row
        )

        fact[
            "fact_type"
        ] = "person_thesis"

        facts.append(
            fact
        )

    for row in work_rows:
        fact = dict(
            row
        )

        fact[
            "fact_type"
        ] = "person_work"

        facts.append(
            fact
        )

    unique_start_rows = (
        deduplicate_rows(
            start_year_rows,
            (
                "person_element_id",
            ),
        )
    )

    for row in unique_start_rows:
        fact = dict(
            row
        )

        fact[
            "fact_type"
        ] = "person_start_year"

        facts.append(
            fact
        )

    unique_person_keywords = (
        deduplicate_rows(
            person_keyword_rows,
            (
                "person_element_id",
                "keyword_element_id",
            ),
        )
    )

    for row in unique_person_keywords:
        fact = dict(
            row
        )

        fact[
            "fact_type"
        ] = "person_keyword"

        facts.append(
            fact
        )

    unique_person_categories = (
        deduplicate_rows(
            person_category_rows,
            (
                "person_element_id",
                "category_element_id",
            ),
        )
    )

    for row in unique_person_categories:
        fact = dict(
            row
        )

        fact[
            "fact_type"
        ] = "person_category"

        facts.append(
            fact
        )

    unique_shared_keywords = (
        deduplicate_rows(
            shared_keyword_rows,
            (
                "left_person_element_id",
                "right_person_element_id",
                "keyword_element_id",
            ),
        )
    )

    for row in unique_shared_keywords:
        fact = dict(
            row
        )

        fact[
            "fact_type"
        ] = "shared_keyword"

        facts.append(
            fact
        )

    unique_shared_categories = (
        deduplicate_rows(
            shared_category_rows,
            (
                "left_person_element_id",
                "right_person_element_id",
                "category_element_id",
            ),
        )
    )

    for row in unique_shared_categories:
        fact = dict(
            row
        )

        fact[
            "fact_type"
        ] = "shared_category"

        facts.append(
            fact
        )

    unique_shared_authors = (
        deduplicate_rows(
            shared_document_author_rows,
            (
                "left_person_element_id",
                "right_person_element_id",
                "document_author_element_id",
            ),
        )
    )

    for row in unique_shared_authors:
        fact = dict(
            row
        )

        fact[
            "fact_type"
        ] = (
            "shared_document_author"
        )

        facts.append(
            fact
        )

    if (
        mentions_shared
        and len(
            unique_start_rows
        )
        >= 2
    ):
        ordered = sorted(
            unique_start_rows,
            key=lambda row: str(
                row.get(
                    "person"
                )
                or ""
            ),
        )

        years = [
            row.get(
                "start_year"
            )
            for row in ordered
        ]

        known_years = [
            year
            for year in years
            if year is not None
        ]

        same = (
            len(
                known_years
            )
            == len(
                years
            )
            and len(
                set(
                    known_years
                )
            )
            == 1
        )

        facts.append(
            {
                "fact_type": (
                    "start_year_comparison"
                ),
                "people": [
                    str(
                        row.get(
                            "person"
                        )
                        or ""
                    )
                    for row in ordered
                ],
                "start_years": {
                    str(
                        row.get(
                            "person"
                        )
                        or ""
                    ): row.get(
                        "start_year"
                    )
                    for row in ordered
                },
                "same_start_year": (
                    same
                ),
                "shared_start_year": (
                    known_years[0]
                    if same
                    else None
                ),
            }
        )

    return facts


def search_graph_with_session(
    session: Any,
    query: str,
    *,
    limit: int = (
        DEFAULT_LIMIT
    ),
    max_entities: int = (
        DEFAULT_MAX_ENTITIES
    ),
) -> dict[str, Any]:
    clean_query = str(
        query
        or ""
    ).strip()

    if not clean_query:
        raise ValueError(
            "Query nesmie byť prázdny"
        )

    if limit < 1:
        raise ValueError(
            "Limit musí byť aspoň 1"
        )

    if max_entities < 1:
        raise ValueError(
            "max_entities musí byť aspoň 1"
        )

    constraints = (
        extract_query_constraints(
            clean_query
        )
    )

    catalog = run_rows(
        session,
        CATALOG_QUERY,
        scope=GRAPH_SCOPE,
    )

    resolved_entities = (
        resolve_entities(
            clean_query,
            catalog,
            max_entities=(
                max_entities
            ),
            preferred_kinds=(
                constraints[
                    "target_kinds"
                ]
            ),
        )
    )

    start_ids = [
        str(
            entity[
                "element_id"
            ]
        )
        for entity
        in resolved_entities
        if entity.get(
            "element_id"
        )
    ]

    person_ids = [
        str(
            entity[
                "element_id"
            ]
        )
        for entity
        in resolved_entities
        if (
            entity.get(
                "element_id"
            )
            and entity.get(
                "kind"
            )
            == "Person"
        )
    ]

    thesis_ids = [
        str(
            entity[
                "element_id"
            ]
        )
        for entity
        in resolved_entities
        if (
            entity.get(
                "element_id"
            )
            and entity.get(
                "kind"
            )
            == "Thesis"
        )
    ]

    work_ids = [
        str(
            entity[
                "element_id"
            ]
        )
        for entity
        in resolved_entities
        if (
            entity.get(
                "element_id"
            )
            and entity.get(
                "kind"
            )
            == "Work"
        )
    ]

    keyword_ids = [
        str(
            entity[
                "element_id"
            ]
        )
        for entity
        in resolved_entities
        if (
            entity.get(
                "element_id"
            )
            and entity.get(
                "kind"
            )
            == "Keyword"
        )
    ]

    category_ids = [
        str(
            entity[
                "element_id"
            ]
        )
        for entity
        in resolved_entities
        if (
            entity.get(
                "element_id"
            )
            and entity.get(
                "kind"
            )
            == "Category"
        )
    ]

    year_specific = bool(
        constraints[
            "years"
        ]
        and not constraints[
            "mentions_start_year"
        ]
    )

    title_specific = (
        constraints[
            "mentions_title"
        ]
    )

    thesis_specific = bool(
        constraints[
            "mentions_thesis"
        ]
        or thesis_ids
        or (
            title_specific
            and not constraints[
                "mentions_work"
            ]
        )
        or (
            year_specific
            and not constraints[
                "mentions_work"
            ]
        )
    )

    work_specific = bool(
        constraints[
            "mentions_work"
        ]
        or work_ids
        or (
            title_specific
            and not constraints[
                "mentions_thesis"
            ]
        )
        or (
            year_specific
            and not constraints[
                "mentions_thesis"
            ]
        )
    )

    shared_keyword_query = bool(
        constraints[
            "mentions_shared"
        ]
        and len(
            person_ids
        )
        >= 2
        and (
            constraints[
                "mentions_keyword"
            ]
            or constraints[
                "mentions_topic"
            ]
        )
    )

    shared_category_query = bool(
        constraints[
            "mentions_shared"
        ]
        and len(
            person_ids
        )
        >= 2
        and constraints[
            "mentions_category"
        ]
    )

    shared_document_author_query = (
        bool(
            constraints[
                "mentions_shared"
            ]
            and len(
                person_ids
            )
            >= 2
            and constraints[
                "mentions_document_author"
            ]
        )
    )

    keyword_specific = bool(
        not constraints[
            "mentions_shared"
        ]
        and (
            keyword_ids
            or (
                person_ids
                and constraints[
                    "mentions_keyword"
                ]
            )
        )
    )

    category_specific = bool(
        not constraints[
            "mentions_shared"
        ]
        and (
            category_ids
            or (
                person_ids
                and constraints[
                    "mentions_category"
                ]
            )
        )
    )

    structured_query = (
        thesis_specific
        or work_specific
        or constraints[
            "mentions_start_year"
        ]
        or keyword_specific
        or category_specific
        or shared_keyword_query
        or shared_category_query
        or shared_document_author_query
    )

    paths: list[
        dict[str, Any]
    ] = []

    if thesis_specific:
        paths.extend(
            run_rows(
                session,
                THESIS_PATHS_QUERY,
                scope=GRAPH_SCOPE,
                person_ids=(
                    person_ids
                ),
                thesis_ids=(
                    thesis_ids
                ),
                years=(
                    constraints[
                        "years"
                    ]
                ),
                thesis_types=(
                    constraints[
                        "thesis_types"
                    ]
                ),
                limit=limit,
            )
        )

    if work_specific:
        paths.extend(
            run_rows(
                session,
                WORK_PATHS_QUERY,
                scope=GRAPH_SCOPE,
                person_ids=(
                    person_ids
                ),
                work_ids=(
                    work_ids
                ),
                years=(
                    constraints[
                        "years"
                    ]
                ),
                work_types=(
                    constraints[
                        "work_types"
                    ]
                ),
                limit=limit,
            )
        )

    if (
        not structured_query
        and start_ids
    ):
        if (
            constraints[
                "mentions_shared"
            ]
            and len(
                start_ids
            )
            >= 2
        ):
            paths.extend(
                run_rows(
                    session,
                    PAIR_PATH_QUERY,
                    scope=GRAPH_SCOPE,
                    start_ids=(
                        start_ids
                    ),
                    limit=limit,
                )
            )

        else:
            paths.extend(
                run_rows(
                    session,
                    NEIGHBORHOOD_QUERY,
                    scope=GRAPH_SCOPE,
                    start_ids=(
                        start_ids
                    ),
                    target_kinds=(
                        constraints[
                            "target_kinds"
                        ]
                    ),
                    limit=limit,
                )
            )

    thesis_rows: list[
        dict[str, Any]
    ] = []

    if thesis_specific:
        thesis_rows = run_rows(
            session,
            THESIS_FACTS_QUERY,
            scope=GRAPH_SCOPE,
            person_ids=(
                person_ids
            ),
            thesis_ids=(
                thesis_ids
            ),
            years=(
                constraints[
                    "years"
                ]
            ),
            thesis_types=(
                constraints[
                    "thesis_types"
                ]
            ),
            limit=limit,
        )

    work_rows: list[
        dict[str, Any]
    ] = []

    if work_specific:
        work_rows = run_rows(
            session,
            WORK_FACTS_QUERY,
            scope=GRAPH_SCOPE,
            person_ids=(
                person_ids
            ),
            work_ids=(
                work_ids
            ),
            years=(
                constraints[
                    "years"
                ]
            ),
            work_types=(
                constraints[
                    "work_types"
                ]
            ),
            limit=limit,
        )

    start_year_rows: list[
        dict[str, Any]
    ] = []

    if constraints[
        "mentions_start_year"
    ]:
        start_year_rows = (
            run_rows(
                session,
                START_YEAR_FACTS_QUERY,
                scope=GRAPH_SCOPE,
                person_ids=(
                    person_ids
                ),
                years=(
                    constraints[
                        "years"
                    ]
                ),
                limit=limit,
            )
        )

    person_keyword_rows: list[
        dict[str, Any]
    ] = []

    if keyword_specific:
        person_keyword_rows = (
            run_rows(
                session,
                PERSON_KEYWORD_FACTS_QUERY,
                scope=GRAPH_SCOPE,
                person_ids=(
                    person_ids
                ),
                keyword_ids=(
                    keyword_ids
                ),
                limit=limit,
            )
        )

    person_category_rows: list[
        dict[str, Any]
    ] = []

    if category_specific:
        person_category_rows = (
            run_rows(
                session,
                PERSON_CATEGORY_FACTS_QUERY,
                scope=GRAPH_SCOPE,
                person_ids=(
                    person_ids
                ),
                category_ids=(
                    category_ids
                ),
                limit=limit,
            )
        )

    shared_keyword_rows: list[
        dict[str, Any]
    ] = []

    if shared_keyword_query:
        shared_keyword_rows = (
            run_rows(
                session,
                SHARED_KEYWORD_FACTS_QUERY,
                scope=GRAPH_SCOPE,
                person_ids=(
                    person_ids
                ),
                limit=limit,
            )
        )

    shared_category_rows: list[
        dict[str, Any]
    ] = []

    if shared_category_query:
        shared_category_rows = (
            run_rows(
                session,
                SHARED_CATEGORY_FACTS_QUERY,
                scope=GRAPH_SCOPE,
                person_ids=(
                    person_ids
                ),
                limit=limit,
            )
        )

    shared_document_author_rows: list[
        dict[str, Any]
    ] = []

    if shared_document_author_query:
        shared_document_author_rows = (
            run_rows(
                session,
                SHARED_DOCUMENT_AUTHOR_FACTS_QUERY,
                scope=GRAPH_SCOPE,
                person_ids=(
                    person_ids
                ),
                limit=limit,
            )
        )

    paths = deduplicate_paths(
        paths
    )

    facts = build_facts(
        thesis_rows,
        work_rows,
        start_year_rows,
        person_keyword_rows,
        person_category_rows,
        shared_keyword_rows,
        shared_category_rows,
        shared_document_author_rows,
        mentions_shared=(
            constraints[
                "mentions_shared"
            ]
        ),
    )

    documents = collect_documents(
        resolved_entities,
        paths,
        facts,
    )

    strategies: list[
        str
    ] = []

    if resolved_entities:
        strategies.append(
            "entity_resolution"
        )

    if paths:
        strategies.append(
            "graph_traversal"
        )

    if thesis_rows:
        strategies.append(
            "thesis_constraints"
        )

    if work_rows:
        strategies.append(
            "work_constraints"
        )

    if start_year_rows:
        strategies.append(
            "start_year_constraints"
        )

    if person_keyword_rows:
        strategies.append(
            "person_keyword_constraints"
        )

    if person_category_rows:
        strategies.append(
            "person_category_constraints"
        )

    if shared_keyword_rows:
        strategies.append(
            "shared_keyword_constraints"
        )

    if shared_category_rows:
        strategies.append(
            "shared_category_constraints"
        )

    if shared_document_author_rows:
        strategies.append(
            "shared_document_author_constraints"
        )

    return {
        "query": clean_query,
        "strategies": strategies,
        "constraints": constraints,
        "resolved_entities": (
            resolved_entities
        ),
        "paths": paths[:limit],
        "facts": facts[:limit],
        "documents": documents,
    }


def search_graph(
    query: str,
    *,
    limit: int = (
        DEFAULT_LIMIT
    ),
    max_entities: int = (
        DEFAULT_MAX_ENTITIES
    ),
) -> dict[str, Any]:
    from .graph_db import (
        create_neo4j_driver,
        get_neo4j_settings,
    )

    settings = (
        get_neo4j_settings()
    )

    with create_neo4j_driver(
        settings
    ) as driver:
        driver.verify_connectivity()

        with driver.session(
            database=(
                settings.database
            )
        ) as session:
            return (
                search_graph_with_session(
                    session,
                    query,
                    limit=limit,
                    max_entities=(
                        max_entities
                    ),
                )
            )


def main() -> None:
    parser = (
        argparse.ArgumentParser()
    )

    parser.add_argument(
        "query"
    )

    parser.add_argument(
        "--limit",
        type=int,
        default=DEFAULT_LIMIT,
    )

    parser.add_argument(
        "--max-entities",
        type=int,
        default=(
            DEFAULT_MAX_ENTITIES
        ),
    )

    args = parser.parse_args()

    result = search_graph(
        args.query,
        limit=args.limit,
        max_entities=(
            args.max_entities
        ),
    )

    print(
        json.dumps(
            result,
            ensure_ascii=False,
            indent=2,
        )
    )


if __name__ == "__main__":
    main()