from __future__ import annotations

import json
import math
import re
import statistics
import unicodedata
from collections import defaultdict
from pathlib import Path
from typing import Any


VALID_SPLITS = (
    "dev",
    "test",
)

FACT_RELATIONS = {
    "person_thesis": {
        "AUTHORED",
        "EVIDENCED_BY",
    },
    "person_work": {
        "HAS_WORK",
        "EVIDENCED_BY",
    },
    "person_keyword": {
        "HAS_DOCUMENT",
        "HAS_KEYWORD",
    },
    "shared_keyword": {
        "HAS_DOCUMENT",
        "HAS_KEYWORD",
    },
    "person_category": {
        "HAS_DOCUMENT",
        "IN_CATEGORY",
    },
    "shared_category": {
        "HAS_DOCUMENT",
        "IN_CATEGORY",
    },
    "person_document_author": {
        "HAS_DOCUMENT",
        "AUTHORED_DOCUMENT",
    },
    "shared_document_author": {
        "HAS_DOCUMENT",
        "AUTHORED_DOCUMENT",
    },
}


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

    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 load_questions(
    path: Path,
) -> list[dict[str, Any]]:
    if not path.exists():
        raise FileNotFoundError(
            f"Benchmark neexistuje: {path}"
        )

    data = json.loads(
        path.read_text(
            encoding="utf-8"
        )
    )

    if not isinstance(
        data,
        list,
    ):
        raise ValueError(
            "GraphRAG benchmark musí byť JSON pole."
        )

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

    seen_ids: set[str] = set()

    for index, item in enumerate(
        data,
        start=1,
    ):
        if not isinstance(
            item,
            dict,
        ):
            raise ValueError(
                f"Položka {index} nie je JSON objekt."
            )

        question_id = item.get(
            "id"
        )

        question = item.get(
            "question"
        )

        split = item.get(
            "split"
        )

        graph = item.get(
            "graph"
        )

        if (
            not isinstance(
                question_id,
                str,
            )
            or not question_id.strip()
        ):
            raise ValueError(
                f"Položka {index} nemá platné id."
            )

        if question_id in seen_ids:
            raise ValueError(
                f"Duplicitné id: {question_id}"
            )

        seen_ids.add(
            question_id
        )

        if (
            not isinstance(
                question,
                str,
            )
            or not question.strip()
        ):
            raise ValueError(
                f"{question_id}: chýba question."
            )

        if split not in VALID_SPLITS:
            raise ValueError(
                f"{question_id}: neplatný split "
                f"{split!r}."
            )

        if not isinstance(
            graph,
            dict,
        ):
            raise ValueError(
                f"{question_id}: chýba graph."
            )

        if graph.get(
            "schema_version"
        ) != 2:
            raise ValueError(
                f"{question_id}: benchmark nie je "
                "schema_version=2."
            )

        questions.append(
            item
        )

    return questions


def filter_questions(
    questions: list[
        dict[str, Any]
    ],
    *,
    split: str,
    ids: set[str] | None = None,
    max_questions: int | None = None,
) -> list[dict[str, Any]]:
    selected = [
        item
        for item in questions
        if (
            split == "all"
            or item.get("split")
            == split
        )
    ]

    if ids:
        selected = [
            item
            for item in selected
            if str(
                item.get("id")
            )
            in ids
        ]

    if max_questions is not None:
        selected = selected[
            :max_questions
        ]

    return selected


def evidence_payload(
    result: dict[str, Any],
) -> dict[str, Any]:
    return {
        "paths": result.get(
            "paths",
            [],
        ),
        "facts": result.get(
            "facts",
            [],
        ),
        "aggregation": result.get("aggregation"),
    }


def evidence_text(
    result: dict[str, Any],
) -> str:
    payload = evidence_payload(
        result
    )

    rendered = json.dumps(
        payload,
        ensure_ascii=False,
        sort_keys=True,
        default=str,
    )

    return normalize_text(
        rendered
    )


def value_is_present(
    value: Any,
    result: dict[str, Any],
) -> bool:
    expected = normalize_text(
        value
    )

    if not expected:
        return False

    actual = evidence_text(
        result
    )

    return f" {expected} " in f" {actual} "


def ratio(
    matched: int,
    expected: int,
) -> float:
    if expected <= 0:
        return 1.0

    return matched / expected


def entity_value_is_present(entity_type: str, value: str, result: dict[str, Any]) -> bool:
    """Score typed answer evidence, never aliases copied from the question."""
    rows = list(result.get("facts", []))
    rows.extend(node for path in result.get("paths", []) for node in path.get("nodes", []))
    fields = {
        "person": ("person", "left_person", "right_person", "people"),
        "keyword": ("keyword",),
        "category": ("category",),
        "document_author": ("document_author",),
        "year": ("year", "start_year", "shared_start_year", "shared_years"),
        "start_year": ("start_year", "shared_start_year"),
        "thesis_title": ("title",),
        "work_title": ("title",),
    }.get(entity_type, ())
    node_kinds = {"person": "Person", "keyword": "Keyword", "category": "Category",
                  "document_author": "DocumentAuthor"}
    expected = normalize_text(value)
    for row in rows:
        if entity_type == "thesis_title" and not (row.get("kind") == "Thesis" or row.get("thesis_type")):
            continue
        if entity_type == "work_title" and not (row.get("kind") == "Work" or row.get("work_type")):
            continue
        candidates = [row.get(field) for field in fields]
        if row.get("kind") == node_kinds.get(entity_type) and row.get("kind") is not None:
            candidates.append(row.get("name"))
        for candidate in candidates:
            values = candidate if isinstance(candidate, list) else [candidate]
            if any(normalize_text(item) == expected for item in values if item is not None):
                return True
    return False


def expected_entity_values(
    question: dict[str, Any],
) -> list[
    tuple[str, str]
]:
    graph = question.get(
        "graph",
        {}
    )

    result: list[
        tuple[str, str]
    ] = []

    for entity in graph.get(
        "expected_entities",
        [],
    ):
        if not isinstance(
            entity,
            dict,
        ):
            continue

        entity_type = str(
            entity.get(
                "type"
            )
            or ""
        )

        value = entity.get(
            "value"
        )

        if (
            entity_type
            and value is not None
        ):
            result.append(
                (
                    entity_type,
                    str(value),
                )
            )

    return result


def collect_actual_relationships(
    result: dict[str, Any],
) -> set[str]:
    relationships: set[str] = set()

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

        for relationship in path.get(
            "relationships",
            [],
        ):
            if relationship:
                relationships.add(
                    str(relationship)
                )

    for fact in result.get(
        "facts",
        [],
    ):
        if not isinstance(
            fact,
            dict,
        ):
            continue

        fact_type = str(
            fact.get(
                "fact_type"
            )
            or ""
        )

        relationships.update(
            FACT_RELATIONS.get(
                fact_type,
                set(),
            )
        )

    return relationships


def property_present(
    property_name: str,
    result: dict[str, Any],
) -> bool:
    if "|" in property_name:
        return any(
            property_present(
                part,
                result,
            )
            for part in property_name.split(
                "|"
            )
        )

    property_name = property_name.strip()

    facts = result.get(
        "facts",
        [],
    )

    path_nodes = []

    for path in result.get(
        "paths",
        [],
    ):
        if isinstance(
            path,
            dict,
        ):
            path_nodes.extend(
                path.get(
                    "nodes",
                    [],
                )
            )

    rows = [
        row
        for row in (
            list(facts)
            + list(path_nodes)
        )
        if isinstance(
            row,
            dict,
        )
    ]

    if property_name == "Person.start_year":
        return any(
            row.get(
                "start_year"
            )
            is not None
            for row in rows
        )

    if property_name == "Thesis.title":
        return any(
            (
                row.get("kind")
                == "Thesis"
                or row.get("thesis_type")
                is not None
                or row.get("fact_type")
                == "person_thesis"
            )
            and bool(
                row.get(
                    "title"
                )
            )
            for row in rows
        )

    if property_name == "Thesis.year":
        return any(
            (
                row.get("kind")
                == "Thesis"
                or row.get("thesis_type")
                is not None
                or row.get("fact_type")
                == "person_thesis"
            )
            and row.get(
                "year"
            )
            is not None
            for row in rows
        )

    if property_name == "Work.title":
        return any(
            (
                row.get("kind")
                == "Work"
                or row.get("work_type")
                is not None
                or row.get("fact_type")
                == "person_work"
            )
            and bool(
                row.get(
                    "title"
                )
            )
            for row in rows
        )

    if property_name == "Work.year":
        return any(
            (
                row.get("kind")
                == "Work"
                or row.get("work_type")
                is not None
                or row.get("fact_type")
                == "person_work"
            )
            and row.get(
                "year"
            )
            is not None
            for row in rows
        )

    return False


def collect_document_paths(
    result: dict[str, Any],
) -> set[str]:
    paths: set[str] = set()

    for document in result.get(
        "documents",
        [],
    ):
        if not isinstance(
            document,
            dict,
        ):
            continue

        path = document.get(
            "document_path"
        )

        if path:
            paths.add(
                str(path)
            )

    return paths


def collect_source_urls(
    result: dict[str, Any],
) -> set[str]:
    urls: set[str] = set()

    def add(
        value: Any,
    ) -> None:
        if value:
            urls.add(
                str(value)
            )

    for document in result.get(
        "documents",
        [],
    ):
        if isinstance(
            document,
            dict,
        ):
            add(
                document.get(
                    "source_url"
                )
            )

    for fact in result.get(
        "facts",
        [],
    ):
        if not isinstance(
            fact,
            dict,
        ):
            continue

        for key, value in fact.items():
            if key.endswith(
                "source_url"
            ):
                add(
                    value
                )

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

        for node in path.get(
            "nodes",
            [],
        ):
            if isinstance(
                node,
                dict,
            ):
                add(
                    node.get(
                        "source_url"
                    )
                )

    return urls


def normalized_set(
    values: list[Any] | set[Any],
) -> set[str]:
    return {
        normalize_text(
            value
        )
        for value in values
        if normalize_text(
            value
        )
    }


def actual_entity_set(
    question: dict[str, Any],
    result: dict[str, Any],
) -> set[str]:
    category = str(
        question.get(
            "category"
        )
        or ""
    )

    values: set[str] = set()

    if category == "person_multiple_categories":
        for fact in result.get(
            "facts",
            [],
        ):
            if (
                isinstance(
                    fact,
                    dict,
                )
                and fact.get(
                    "fact_type"
                )
                == "person_category"
                and fact.get(
                    "category"
                )
            ):
                values.add(
                    str(
                        fact[
                            "category"
                        ]
                    )
                )

        return values

    if category == "candidate_set_aggregation":
        for fact in result.get(
            "facts",
            [],
        ):
            if not isinstance(
                fact,
                dict,
            ):
                continue

            if fact.get(
                "fact_type"
            ) in {
                "person_thesis",
                "person_work",
                "person_academic_year",
            }:
                person = fact.get(
                    "person"
                )

                if person:
                    values.add(
                        str(person)
                    )

        return values

    return values


def entity_set_metrics(
    question: dict[str, Any],
    result: dict[str, Any],
) -> dict[str, Any]:
    graph = question.get(
        "graph",
        {}
    )

    expected = graph.get(
        "expected_entity_set"
    )

    if not isinstance(
        expected,
        list,
    ):
        return {
            "entity_set_expected": [],
            "entity_set_actual": [],
            "entity_set_mode": None,
            "entity_set_ok": True,
        }

    actual = actual_entity_set(
        question,
        result,
    )

    expected_normalized = (
        normalized_set(
            expected
        )
    )

    actual_normalized = (
        normalized_set(
            actual
        )
    )

    mode = str(
        graph.get(
            "entity_set_mode"
        )
        or "contains"
    )

    if mode == "exact":
        ok = (
            actual_normalized
            == expected_normalized
        )
    else:
        ok = (
            expected_normalized
            <= actual_normalized
        )

    return {
        "entity_set_expected": (
            sorted(
                expected_normalized
            )
        ),
        "entity_set_actual": (
            sorted(
                actual_normalized
            )
        ),
        "entity_set_mode": mode,
        "entity_set_ok": ok,
    }


def evaluate_result(
    question: dict[str, Any],
    result: dict[str, Any],
    *,
    latency_seconds: float,
    error: str | None = None,
) -> dict[str, Any]:
    graph = question.get(
        "graph",
        {}
    )

    expected_values = [
        str(value)
        for value in question.get(
            "expected_answer_contains",
            [],
        )
    ]

    matched_values = [
        value
        for value in expected_values
        if value_is_present(
            value,
            result,
        )
    ]

    entities = expected_entity_values(
        question
    )

    matched_entities = [
        {
            "type": entity_type,
            "value": value,
        }
        for entity_type, value
        in entities
        if entity_value_is_present(
            entity_type,
            value,
            result,
        )
    ]

    expected_relations = [
        str(value)
        for value in graph.get(
            "expected_relations",
            [],
        )
    ]

    actual_relations = (
        collect_actual_relationships(
            result
        )
    )

    matched_relations = [
        relationship
        for relationship
        in expected_relations
        if relationship
        in actual_relations
    ]

    expected_properties = [
        str(value)
        for value in graph.get(
            "expected_properties",
            [],
        )
    ]

    matched_properties = [
        value
        for value
        in expected_properties
        if property_present(
            value,
            result,
        )
    ]

    relation_choices = graph.get(
        "expected_relation_choices",
        [],
    )

    relation_choice_results = []

    if isinstance(
        relation_choices,
        list,
    ):
        for choice in relation_choices:
            values = (
                choice
                if isinstance(
                    choice,
                    list,
                )
                else [choice]
            )

            clean = [
                str(value)
                for value in values
                if value
            ]

            relation_choice_results.append(
                {
                    "choices": clean,
                    "matched": any(
                        value
                        in actual_relations
                        for value in clean
                    ),
                }
            )

    expected_documents = {
        str(value)
        for value in question.get(
            "expected_documents",
            [],
        )
    }

    actual_documents = (
        collect_document_paths(
            result
        )
    )

    matched_documents = (
        expected_documents
        & actual_documents
    )

    expected_sources = {
        str(value)
        for value in question.get(
            "expected_source_urls",
            [],
        )
    }

    actual_sources = (
        collect_source_urls(
            result
        )
    )

    matched_sources = (
        expected_sources
        & actual_sources
    )

    set_metrics = entity_set_metrics(
        question,
        result,
    )

    value_score = ratio(
        len(matched_values),
        len(expected_values),
    )

    entity_score = ratio(
        len(matched_entities),
        len(entities),
    )

    relation_score = ratio(
        len(matched_relations),
        len(expected_relations),
    )

    property_score = ratio(
        len(matched_properties),
        len(expected_properties),
    )

    relation_choice_score = ratio(
        sum(
            1
            for item
            in relation_choice_results
            if item["matched"]
        ),
        len(
            relation_choice_results
        ),
    )

    document_recall = ratio(
        len(matched_documents),
        len(expected_documents),
    )

    source_recall = ratio(
        len(matched_sources),
        len(expected_sources),
    )

    expected_count = graph.get("expected_count")
    aggregation = result.get("aggregation") or {}
    count_ok = expected_count is None or (
        aggregation.get("complete") is True
        and aggregation.get("count") == expected_count
    )
    complete = not result.get("truncated", False)
    strict = all(
        (
            error is None,
            complete,
            count_ok,
            value_score == 1.0,
            entity_score == 1.0,
            relation_score == 1.0,
            property_score == 1.0,
            relation_choice_score
            == 1.0,
            document_recall == 1.0,
            source_recall == 1.0,
            bool(
                set_metrics[
                    "entity_set_ok"
                ]
            ),
        )
    )

    return {
        "count_ok": count_ok,
        "expected_count": expected_count,
        "actual_count": aggregation.get("count"),
        "complete": complete,
        "id": question.get(
            "id"
        ),
        "split": question.get(
            "split"
        ),
        "category": question.get(
            "category"
        ),
        "difficulty": question.get(
            "difficulty"
        ),
        "question": question.get(
            "question"
        ),
        "latency_seconds": round(
            latency_seconds,
            6,
        ),
        "error": error,
        "strict": strict,
        "value_score": round(
            value_score,
            6,
        ),
        "entity_score": round(
            entity_score,
            6,
        ),
        "relation_score": round(
            relation_score,
            6,
        ),
        "property_score": round(
            property_score,
            6,
        ),
        "relation_choice_score": round(
            relation_choice_score,
            6,
        ),
        "document_recall": round(
            document_recall,
            6,
        ),
        "source_recall": round(
            source_recall,
            6,
        ),
        "expected_values": (
            expected_values
        ),
        "matched_values": (
            matched_values
        ),
        "expected_entities": [
            {
                "type": entity_type,
                "value": value,
            }
            for entity_type, value
            in entities
        ],
        "matched_entities": (
            matched_entities
        ),
        "expected_relations": (
            expected_relations
        ),
        "actual_relations": sorted(
            actual_relations
        ),
        "matched_relations": (
            matched_relations
        ),
        "expected_properties": (
            expected_properties
        ),
        "matched_properties": (
            matched_properties
        ),
        "relation_choices": (
            relation_choice_results
        ),
        "expected_documents": sorted(
            expected_documents
        ),
        "actual_documents": sorted(
            actual_documents
        ),
        "matched_documents": sorted(
            matched_documents
        ),
        "expected_source_urls": sorted(
            expected_sources
        ),
        "actual_source_urls": sorted(
            actual_sources
        ),
        "matched_source_urls": sorted(
            matched_sources
        ),
        **set_metrics,
        "strategies": list(
            result.get(
                "strategies",
                [],
            )
        ),
        "fact_types": sorted(
            {
                str(
                    fact.get(
                        "fact_type"
                    )
                )
                for fact in result.get(
                    "facts",
                    [],
                )
                if (
                    isinstance(
                        fact,
                        dict,
                    )
                    and fact.get(
                        "fact_type"
                    )
                )
            }
        ),
    }


def percentile(
    values: list[float],
    fraction: float,
) -> float:
    if not values:
        return 0.0

    ordered = sorted(
        values
    )

    if len(ordered) == 1:
        return ordered[0]

    position = (
        (len(ordered) - 1)
        * fraction
    )

    lower = math.floor(
        position
    )

    upper = math.ceil(
        position
    )

    if lower == upper:
        return ordered[
            lower
        ]

    weight = (
        position - lower
    )

    return (
        ordered[lower]
        * (1.0 - weight)
        + ordered[upper]
        * weight
    )


def aggregate_rows(
    rows: list[
        dict[str, Any]
    ],
) -> dict[str, Any]:
    if not rows:
        return {
            "questions": 0,
            "strict_count": 0,
            "strict_accuracy": 0.0,
            "error_count": 0,
            "value_score": 0.0,
            "entity_score": 0.0,
            "relation_score": 0.0,
            "property_score": 0.0,
            "relation_choice_score": 0.0,
            "document_recall": 0.0,
            "source_recall": 0.0,
            "latency_mean_seconds": 0.0,
            "latency_median_seconds": 0.0,
            "latency_p95_seconds": 0.0,
        }

    latencies = [
        float(
            row[
                "latency_seconds"
            ]
        )
        for row in rows
    ]

    def mean(
        key: str,
    ) -> float:
        return round(
            statistics.fmean(
                float(
                    row[key]
                )
                for row in rows
            ),
            6,
        )

    strict_count = sum(
        bool(
            row["strict"]
        )
        for row in rows
    )

    return {
        "questions": len(
            rows
        ),
        "strict_count": (
            strict_count
        ),
        "strict_accuracy": round(
            strict_count
            / len(rows),
            6,
        ),
        "error_count": sum(
            bool(
                row.get(
                    "error"
                )
            )
            for row in rows
        ),
        "value_score": mean(
            "value_score"
        ),
        "entity_score": mean(
            "entity_score"
        ),
        "relation_score": mean(
            "relation_score"
        ),
        "property_score": mean(
            "property_score"
        ),
        "relation_choice_score": (
            mean(
                "relation_choice_score"
            )
        ),
        "document_recall": mean(
            "document_recall"
        ),
        "source_recall": mean(
            "source_recall"
        ),
        "latency_mean_seconds": round(
            statistics.fmean(
                latencies
            ),
            6,
        ),
        "latency_median_seconds": round(
            statistics.median(
                latencies
            ),
            6,
        ),
        "latency_p95_seconds": round(
            percentile(
                latencies,
                0.95,
            ),
            6,
        ),
    }


def aggregate_by_field(
    rows: list[
        dict[str, Any]
    ],
    field: str,
) -> dict[
    str,
    dict[str, Any],
]:
    grouped: dict[
        str,
        list[dict[str, Any]],
    ] = defaultdict(list)

    for row in rows:
        key = str(
            row.get(
                field
            )
            or "<missing>"
        )

        grouped[
            key
        ].append(
            row
        )

    return {
        key: aggregate_rows(
            group_rows
        )
        for key, group_rows
        in sorted(
            grouped.items()
        )
    }