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_TOPIC|"
    "DESCRIBES|HAS_WORK|EVIDENCED_BY"
)

CATALOG_QUERY = """
MATCH (n)
WHERE n.graph_scope = $scope
  AND (
    n:Person
    OR n:Topic
    OR n:Category
    OR n:Author
    OR n:Work
  )
RETURN
  elementId(n) AS element_id,
  CASE
    WHEN n:Person THEN 'Person'
    WHEN n:Topic THEN 'Topic'
    WHEN n:Category THEN 'Category'
    WHEN n:Author THEN 'Author'
    WHEN n:Work THEN 'Work'
    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.work_type AS work_type,
  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
        ('Topic' IN $target_kinds AND node:Topic)
        OR ('Category' IN $target_kinds AND node:Category)
        OR ('Author' IN $target_kinds AND node:Author)
        OR ('Work' IN $target_kinds AND node:Work)
        OR ('Document' IN $target_kinds AND node:Document)
        OR ('Person' IN $target_kinds AND node:Person)
    )
  )
RETURN
  [node IN nodes(p) | {{
    element_id: elementId(node),
    kind: CASE
      WHEN node:Person THEN 'Person'
      WHEN node:Document THEN 'Document'
      WHEN node:Work THEN 'Work'
      WHEN node:Topic THEN 'Topic'
      WHEN node:Category THEN 'Category'
      WHEN node:Author THEN 'Author'
      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,
    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)
  AND (NOT $require_topic OR any(node IN nodes(p) WHERE node:Topic))
  AND (NOT $require_category OR any(node IN nodes(p) WHERE node:Category))
  AND (NOT $require_author OR any(node IN nodes(p) WHERE node:Author))
RETURN
  [node IN nodes(p) | {{
    element_id: elementId(node),
    kind: CASE
      WHEN node:Person THEN 'Person'
      WHEN node:Document THEN 'Document'
      WHEN node:Work THEN 'Work'
      WHEN node:Topic THEN 'Topic'
      WHEN node:Category THEN 'Category'
      WHEN node:Author THEN 'Author'
      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,
    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
"""

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
"""

START_YEAR_FACTS_QUERY = """
MATCH (person:Person)-[:HAS_DOCUMENT]->(document:Document)
WHERE person.graph_scope = $scope
  AND document.graph_scope = $scope
  AND (
    size($person_ids) = 0
    OR elementId(person) IN $person_ids
  )
  AND (
    size($years) = 0
    OR person.start_year IN $years
  )
RETURN
  elementId(person) AS person_element_id,
  person.name AS person,
  person.start_year AS start_year,
  document.path AS document_path,
  document.source_url AS source_url
ORDER BY person
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:Document THEN 'Document'
      WHEN node:Work THEN 'Work'
      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,
    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
"""

SHARED_TOPIC_FACTS_QUERY = """
MATCH (left:Person)-[:HAS_DOCUMENT]->(left_document:Document)-[:HAS_TOPIC]->(topic:Topic)
MATCH (right:Person)-[:HAS_DOCUMENT]->(right_document:Document)-[:HAS_TOPIC]->(topic)
WHERE left.graph_scope = $scope
  AND right.graph_scope = $scope
  AND left_document.graph_scope = $scope
  AND right_document.graph_scope = $scope
  AND topic.graph_scope = $scope
  AND elementId(left) IN $person_ids
  AND elementId(right) IN $person_ids
  AND elementId(left) < elementId(right)
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(topic) AS topic_element_id,
  topic.id AS topic_id,
  topic.name AS topic
ORDER BY left_person, right_person, topic
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,
            )
        }
    )

    work_types: list[str] = []

    has_work_word = has_stem(
        tokens,
        "prac",
    )

    has_project_word = has_stem(
        tokens,
        "projekt",
    )

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

    if (
        (
            has_stem(tokens, "diplom")
            and has_work_word
        )
        or "dp" in token_set
    ):
        work_types.append(
            "diploma_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 (
        (
            has_stem(tokens, "dizert")
            and has_work_word
        )
        or "phd" in token_set
    ):
        work_types.append(
            "doctoral_thesis"
        )

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

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

    work_types = list(
        dict.fromkeys(work_types)
    )

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

    mentions_topic = (
        has_stem(tokens, "tem")
        or "tag" in token_set
        or "tagy" in token_set
    )

    mentions_category = has_stem(
        tokens,
        "kategori",
    )

    mentions_author = (
        has_stem(tokens, "autor")
        or has_stem(tokens, "veduc")
    )

    mentions_title = has_stem(
        tokens,
        "nazov",
    )

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

    target_kinds: list[str] = []

    if mentions_topic:
        target_kinds.append("Topic")

    if mentions_category:
        target_kinds.append("Category")

    if mentions_author:
        target_kinds.append("Author")

    if (
        mentions_title
        or work_types
        or (
            years
            and not mentions_start_year
        )
    ):
        target_kinds.append("Work")

    return {
        "years": years,
        "work_types": work_types,
        "mentions_start_year": (
            mentions_start_year
        ),
        "mentions_topic": mentions_topic,
        "mentions_category": (
            mentions_category
        ),
        "mentions_author": mentions_author,
        "mentions_title": mentions_title,
        "mentions_shared": mentions_shared,
        "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",
        "Author",
    }:
        values.append(
            entity.get("name")
        )

    elif kind == "Topic":
        values.extend(
            [
                entity.get("name"),
                entity.get("id"),
            ]
        )

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

    elif kind == "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 = (
        " "
        + normalized_query
        + " "
    )

    padded_alias = (
        " "
        + 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 resolve_entities(
    query: str,
    catalog: list[dict[str, Any]],
    *,
    max_entities: int = DEFAULT_MAX_ENTITIES,
) -> list[dict[str, Any]]:
    normalized_query = normalize_for_match(
        query
    )

    query_tokens = set(
        normalized_query.split()
    )

    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 > 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") == "Author"
            and str(
                item.get("matched_alias")
                or ""
            )
            in person_aliases
        )
    ]

    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 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(
    work_rows: list[dict[str, Any]],
    start_year_rows: list[dict[str, Any]],
    shared_topic_rows: list[dict[str, Any]],
    *,
    mentions_shared: bool,
) -> list[dict[str, Any]]:
    facts: list[dict[str, Any]] = []

    for row in work_rows:
        fact = dict(row)
        fact["fact_type"] = (
            "person_work"
        )
        facts.append(fact)

    for row in start_year_rows:
        fact = dict(row)
        fact["fact_type"] = (
            "person_start_year"
        )
        facts.append(fact)

    for row in shared_topic_rows:
        fact = dict(row)
        fact["fact_type"] = (
            "shared_topic"
        )
        facts.append(fact)

    if (
        mentions_shared
        and len(start_year_rows) >= 2
    ):
        ordered = sorted(
            start_year_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,
    )

    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"
        )
    ]

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

    work_specific = bool(
        constraints["work_types"]
        or constraints["mentions_title"]
        or work_ids
        or (
            constraints["years"]
            and not constraints[
                "mentions_start_year"
            ]
        )
    )

    shared_relation_query = bool(
        constraints["mentions_shared"]
        and len(person_ids) >= 2
        and (
            constraints["mentions_topic"]
            or constraints[
                "mentions_category"
            ]
            or constraints[
                "mentions_author"
            ]
        )
    )

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

    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,
            )
        )

    elif shared_relation_query:
        paths.extend(
            run_rows(
                session,
                PAIR_PATH_QUERY,
                scope=GRAPH_SCOPE,
                start_ids=person_ids,
                require_topic=(
                    constraints[
                        "mentions_topic"
                    ]
                ),
                require_category=(
                    constraints[
                        "mentions_category"
                    ]
                ),
                require_author=(
                    constraints[
                        "mentions_author"
                    ]
                ),
                limit=limit,
            )
        )

    elif not constraints[
        "mentions_start_year"
    ] and start_ids:
        paths.extend(
            run_rows(
                session,
                NEIGHBORHOOD_QUERY,
                scope=GRAPH_SCOPE,
                start_ids=start_ids,
                target_kinds=(
                    constraints[
                        "target_kinds"
                    ]
                ),
                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,
        )

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

    if (
        constraints["mentions_shared"]
        and constraints["mentions_topic"]
        and len(person_ids) >= 2
    ):
        shared_topic_rows = run_rows(
            session,
            SHARED_TOPIC_FACTS_QUERY,
            scope=GRAPH_SCOPE,
            person_ids=person_ids,
            limit=limit,
        )

    paths = deduplicate_paths(
        paths
    )

    facts = build_facts(
        work_rows,
        start_year_rows,
        shared_topic_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 work_rows:
        strategies.append(
            "work_constraints"
        )

    if start_year_rows:
        strategies.append(
            "start_year_constraints"
        )

    if shared_topic_rows:
        strategies.append(
            "shared_topic_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 scripts.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()