from __future__ import annotations

from tools.evaluation.urls import MARKDOWN_URL_RE, normalize_url

import re
import statistics
import unicodedata
from collections import defaultdict
from typing import Any

from tools.answers.context_policy import citation_check

NO_ANSWER_TEXT = (
    "V dostupných dokumentoch ZP Wiki sa túto " "informáciu nepodarilo spoľahlivo nájsť."
)


WORD_RE = re.compile(
    r"[^\W_]+",
    re.UNICODE,
)

UNIT_RE = re.compile(
    r"(?<!\w)(\d+(?:[.,]\d+)?)\s*(kb|mb|gb|tb)(?!\w)",
    re.IGNORECASE,
)

SEMANTIC_STOPWORDS = {
    "a",
    "aj",
    "and",
    "from",
    "of",
    "pomocou",
    "pre",
    "the",
    "v",
    "vo",
    "z",
    "zo",
}

SEMANTIC_TRIGGER_TOKENS = {
    "small",
    "slovak",
    "question",
    "improve",
    "context",
    "answer",
    "accuracy",
    "annotation",
    "cloud",
    "entity",
    "evaluation",
    "extract",
    "hate",
    "hybrid",
    "insert",
    "knowledge",
    "language",
    "llm",
    "medical",
    "medicine",
    "mteb",
    "multiple",
    "multilingual",
    "named",
    "ner",
    "package",
    "recognize",
    "schema",
    "set",
    "training",
    "triplet",
    "unknown",
}


def normalize_text(
    value: str,
) -> str:
    value = UNIT_RE.sub(
        lambda match: (f"{match.group(1)}" f"{match.group(2).lower()}"),
        value,
    )

    value = unicodedata.normalize(
        "NFKC",
        value,
    ).casefold()

    return " ".join(value.split())


def looks_like_no_answer(
    value: str,
) -> bool:
    normalized = normalize_text(value)

    patterns = (
        (r"\bnepodarilo\b" r".{0,180}" r"\b(?:nájsť|zistiť|dohľadať)\b"),
        (
            r"\b(?:nie je|nie sú|nebol|nebola|nebolo|neboli)\b"
            r".{0,180}"
            r"\b(?:uveden|špecifikovan)\w*"
        ),
        (r"\b(?:neobsahuje|neobsahujú)\b" r".{0,180}" r"\b(?:informáci|údaj|špecifikáci)\w*"),
    )

    return any(
        re.search(
            pattern,
            normalized,
            flags=re.DOTALL,
        )
        is not None
        for pattern in patterns
    )


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

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

    return value.casefold()


def match_tokens(
    value: str,
) -> list[str]:
    normalized = unicodedata.normalize(
        "NFKC",
        value,
    )

    return [normalize_match_token(token) for token in WORD_RE.findall(normalized)]


def morphology_token_matches(
    expected: str,
    actual: str,
) -> bool:
    if expected == actual:
        return True

    if expected.isdigit() or actual.isdigit():
        return False

    shorter = min(
        len(expected),
        len(actual),
    )

    if shorter <= 2:
        return False

    required_prefix = 3 if shorter <= 4 else 4 if shorter <= 6 else 5

    return expected[:required_prefix] == actual[:required_prefix]


def _morphology_phrase_matches(
    expected: str,
    answer: str,
) -> bool:
    normalized_expected = normalize_text(expected)

    normalized_answer = normalize_text(answer)

    if not normalized_expected:
        return True

    expected_tokens = match_tokens(normalized_expected)

    answer_tokens = match_tokens(normalized_answer)

    has_numeric_token = any(token.isdigit() for token in expected_tokens)

    if not has_numeric_token and normalized_expected in normalized_answer:
        return True

    if not expected_tokens or len(answer_tokens) < len(expected_tokens):
        return False

    window_size = len(expected_tokens)

    for start in range(len(answer_tokens) - window_size + 1):
        window = answer_tokens[start : start + window_size]

        if all(
            morphology_token_matches(
                expected_token,
                actual_token,
            )
            for (
                expected_token,
                actual_token,
            ) in zip(
                expected_tokens,
                window,
            )
        ):
            return True

    return False


def _semantic_token(
    token: str,
) -> str:
    token = normalize_match_token(token)

    if token == "llm":
        return "llm"

    number_words = {
        "nula": "0",
        "jeden": "1",
        "jedna": "1",
        "jedno": "1",
        "dva": "2",
        "dve": "2",
        "tri": "3",
        "styri": "4",
        "pat": "5",
        "sest": "6",
        "sedem": "7",
        "osem": "8",
        "devat": "9",
        "desat": "10",
        "zero": "0",
        "one": "1",
        "two": "2",
        "three": "3",
        "four": "4",
        "five": "5",
        "six": "6",
        "seven": "7",
        "eight": "8",
        "nine": "9",
        "ten": "10",
    }

    if token in number_words:
        return number_words[token]

    if token.startswith("velk") or token == "large":
        return "large"

    if token.startswith("jazyk") or token == "language" or token == "languages":
        return "language"

    if token.startswith("model"):
        return "model"

    if token.startswith("pomenov") or token.startswith("menovan") or token == "named":
        return "named"

    if token.startswith("entit") or token in {
        "entity",
        "entities",
    }:
        return "entity"

    if token.startswith("anot") or token in {
        "annotation",
        "annotated",
    }:
        return "annotation"

    if token.startswith("znalost") or token == "knowledge":
        return "knowledge"

    if token.startswith("graf") or token == "graph":
        return "graph"

    if token.startswith("hybrid"):
        return "hybrid"

    if token.startswith("cloud") or token.startswith("klaud"):
        return "cloud"

    if (
        token.startswith("multiling")
        or token.startswith("multijaz")
        or token.startswith("viacjazy")
        or (token.startswith("viac") and "jazy" in token)
        or (token.startswith("mnoho") and "jazy" in token)
    ):
        return "multilingual"

    if (
        token.startswith("viacer")
        or token.startswith("mnoh")
        or token.startswith("niekolk")
        or token == "multiple"
    ):
        return "multiple"

    if token.startswith("extrak") or token.startswith("extrah") or token.startswith("extract"):
        return "extract"

    if token.startswith("trojic") or token.startswith("triplet"):
        return "triplet"

    if token.startswith("medicin") or token.startswith("lekars") or token == "medical":
        return "medical"

    if token.startswith("liek") or token in {
        "drug",
        "drugs",
        "medicine",
        "medicines",
        "medication",
        "medications",
    }:
        return "medicine"

    if token.startswith("packag") or token.startswith("balick") or token.startswith("pribal"):
        return "package"

    if token.startswith("insert") or token.startswith("letak"):
        return "insert"

    if (
        token
        in {
            "data",
            "dat",
        }
        or token.startswith("udaj")
        or token.startswith("obsah")
    ):
        return "data"

    if token.startswith("rozpozn") or token.startswith("rozozn") or token == "recognize":
        return "recognize"

    if token.startswith("neznam") or token == "unknown":
        return "unknown"

    if token.startswith("manual") or token == "manually":
        return "manual"

    if token.startswith("trenovac") or token == "training":
        return "training"

    if token.startswith("mnozin") or token.startswith("sad") or token == "set":
        return "set"

    if token.startswith("schem") or token == "schema":
        return "schema"

    if token.startswith("sloven") or token == "slovak":
        return "slovak"

    if token.startswith("zleps") or token == "improve":
        return "improve"

    if token.startswith("presn") or token == "accuracy":
        return "accuracy"

    if (
        token.startswith("kratk")
        or token.startswith("mal")
        or token
        in {
            "short",
            "small",
        }
    ):
        return "small"

    if token.startswith("kontext") or token == "context":
        return "context"

    if token == "hate" or token.startswith("nenavist"):
        return "hate"

    if token.startswith("speech") or token.startswith("prejav") or token.startswith("rec"):
        return "speech"

    if token.startswith("otaz") or token == "question":
        return "question"

    if token.startswith("odpoved") or token.startswith("answer"):
        return "answer"

    if token == "mteb":
        return "mteb"

    if token.startswith("evalu") or token.startswith("hodnot"):
        return "evaluation"

    if token.startswith("sentence"):
        return "sentence"

    if token.startswith("transformer"):
        return "transformer"

    if token == "ner":
        return "ner"

    if token.startswith("korpus") or token == "corpus":
        return "set"

    return token


def semantic_tokens(
    value: str,
) -> list[str]:
    tokens: list[str] = []

    for token in match_tokens(value):
        canonical = _semantic_token(token)

        if canonical in SEMANTIC_STOPWORDS:
            continue

        tokens.append(canonical)

    return tokens


def _contains_llm_concept(
    tokens: set[str],
) -> bool:
    return (
        "llm" in tokens
        or {
            "large",
            "language",
            "model",
        }
        <= tokens
    )


def _contains_multilingual_concept(
    tokens: set[str],
) -> bool:
    return (
        "multilingual" in tokens
        or {
            "multiple",
            "language",
        }
        <= tokens
    )


def _semantic_concept_match(
    expected: str,
    answer: str,
) -> bool:
    expected_tokens = semantic_tokens(expected)

    answer_tokens = semantic_tokens(answer)

    if not expected_tokens or not answer_tokens:
        return False

    expected_set = set(expected_tokens)

    answer_set = set(answer_tokens)

    if expected_set and all(token.isdigit() for token in expected_set):
        return expected_set <= answer_set

    if "llm" in expected_set:
        return _contains_llm_concept(answer_set)

    if {
        "large",
        "language",
        "model",
    } <= expected_set:
        return _contains_llm_concept(answer_set)

    if "ner" in expected_set:
        return (
            "ner" in answer_set
            or {
                "named",
                "entity",
            }
            <= answer_set
        )

    if {
        "named",
        "entity",
    } <= expected_set and "ner" in answer_set:
        return True

    if "multilingual" in expected_set:
        required_tokens = expected_set - {
            "multilingual",
        }

        if _contains_multilingual_concept(answer_set) and required_tokens <= answer_set:
            return True

    if {
        "medical",
        "package",
        "insert",
    } <= expected_set:
        remaining = expected_set - {
            "medical",
            "package",
            "insert",
        }

        if (
            {
                "package",
                "insert",
            }
            <= answer_set
            and (
                {
                    "medical",
                    "medicine",
                }
                & answer_set
            )
            and remaining <= answer_set
        ):
            return True

    if expected_set & SEMANTIC_TRIGGER_TOKENS and expected_set <= answer_set:
        return True

    if "mteb" in expected_set:
        mteb_supported = (
            "mteb" in answer_set
            or {
                "sentence",
                "transformer",
                "evaluation",
            }
            <= answer_set
        )

        hate_supported = (
            not {
                "hate",
                "speech",
            }
            <= expected_set
            or {
                "hate",
                "speech",
            }
            <= answer_set
        )

        if mteb_supported and hate_supported:
            return True

    return False


def expected_phrase_matches(
    expected: str,
    answer: str,
) -> bool:
    if _morphology_phrase_matches(
        expected,
        answer,
    ):
        return True

    compact_expected = normalize_text(expected)

    compact_answer = normalize_text(answer)

    if compact_expected:
        literal_pattern = re.compile(rf"(?<!\w)" rf"{re.escape(compact_expected)}" rf"(?!\w)")

        if literal_pattern.search(compact_answer):
            return True

    return _semantic_concept_match(
        expected,
        answer,
    )


def valid_expected_url(
    value: Any,
) -> str | None:
    text = str(value or "").strip()

    if not text or text in {
        "...",
        "…",
    }:
        return None

    normalized = normalize_url(text)

    if not normalized.startswith(
        (
            "http://",
            "https://",
        )
    ):
        return None

    return normalized


def required_source_match_count(
    question: dict[str, Any],
    expected_urls: list[str],
) -> int:
    if not expected_urls:
        return 0

    if str(question.get("category") or "") == "multi_document":
        return min(
            2,
            len(expected_urls),
        )

    return len(expected_urls)


def evaluate_answer(
    question: dict[str, Any],
    answer: str,
    *,
    tool_called: bool,
    allowed_source_urls: list[str] | None = None,
) -> dict[str, Any]:
    normalized_answer = normalize_text(answer)

    expected_contains = question.get(
        "expected_answer_contains",
        [],
    )

    if not isinstance(
        expected_contains,
        list,
    ):
        expected_contains = []

    answer_matches = [
        expected_phrase_matches(
            str(expected),
            answer,
        )
        for expected in expected_contains
    ]

    answer_contains_score = sum(answer_matches) / len(answer_matches) if answer_matches else 1.0

    raw_expected_urls = question.get(
        "expected_source_urls",
        [],
    )

    if not isinstance(
        raw_expected_urls,
        list,
    ):
        raw_expected_urls = []

    normalized_expected_urls = [
        normalized
        for normalized in (valid_expected_url(url) for url in raw_expected_urls)
        if normalized is not None
    ]

    source_matches = [expected_url in answer for expected_url in normalized_expected_urls]

    source_match_count = sum(source_matches)

    source_required_count = (
        0
        if not bool(
            question.get(
                "should_answer",
                True,
            )
        )
        else required_source_match_count(
            question,
            normalized_expected_urls,
        )
    )

    if source_required_count == 0:
        source_url_score = 1.0
        source_ok = True

    else:
        source_url_score = min(
            1.0,
            source_match_count / source_required_count,
        )

        source_ok = source_match_count >= source_required_count

    should_answer = bool(
        question.get(
            "should_answer",
            True,
        )
    )

    exact_no_answer = normalize_text(NO_ANSWER_TEXT) in normalized_answer

    returned_no_answer = (
        exact_no_answer if should_answer else (exact_no_answer or looks_like_no_answer(answer))
    )

    if should_answer:
        should_answer_ok = bool(answer.strip()) and not returned_no_answer

    else:
        should_answer_ok = returned_no_answer

    tool_score = 1.0 if tool_called else 0.0

    should_answer_score = 1.0 if should_answer_ok else 0.0

    overall_score = statistics.mean(
        [
            answer_contains_score,
            source_url_score,
            should_answer_score,
            tool_score,
        ]
    )

    citations = (
        citation_check({"answer": answer, "source_urls": []}, set(allowed_source_urls))
        if allowed_source_urls is not None
        else {"citations_ok": None, "unsupported_urls": []}
    )
    strict_pass = (
        bool(answer.strip())
        and all(answer_matches)
        and source_ok
        and should_answer_ok
        and tool_called
        and citations["citations_ok"] is not False
    )

    return {
        **citations,
        "citations_checked": allowed_source_urls is not None,
        "answer_matches": (answer_matches),
        "answer_contains_score": (answer_contains_score),
        "source_matches": (source_matches),
        "source_match_count": (source_match_count),
        "source_required_count": (source_required_count),
        "source_url_score": (source_url_score),
        "should_answer_ok": (should_answer_ok),
        "should_answer_score": (should_answer_score),
        "returned_no_answer": (returned_no_answer),
        "tool_score": (tool_score),
        "overall_score": (overall_score),
        "strict_pass": (strict_pass),
    }


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

    return float(statistics.mean(values))


def summarize_results(
    results: list[dict[str, Any]],
    *,
    include_groups: bool = True,
) -> dict[str, Any]:
    total = len(results)

    errors = [item for item in results if item.get("error")]

    completed = total - len(errors)

    tool_called_values = [(1.0 if item.get("tool_called") else 0.0) for item in results]

    answer_scores = [
        float(
            item.get(
                "answer_contains_score",
                0.0,
            )
        )
        for item in results
    ]

    source_scores = [
        float(
            item.get(
                "source_url_score",
                0.0,
            )
        )
        for item in results
    ]

    should_answer_scores = [
        float(
            item.get(
                "should_answer_score",
                0.0,
            )
        )
        for item in results
    ]

    overall_scores = [
        float(
            item.get(
                "overall_score",
                0.0,
            )
        )
        for item in results
    ]

    latencies = [
        float(
            item.get(
                "total_latency_seconds",
                0.0,
            )
        )
        for item in results
        if not item.get("error")
    ]

    strict_passes = sum(1 for item in results if item.get("strict_pass"))

    prompt_tokens = sum(
        int(
            (item.get("usage") or {}).get(
                "prompt_tokens",
                0,
            )
        )
        for item in results
    )

    completion_tokens = sum(
        int(
            (item.get("usage") or {}).get(
                "completion_tokens",
                0,
            )
        )
        for item in results
    )

    total_tokens = sum(
        int(
            (item.get("usage") or {}).get(
                "total_tokens",
                0,
            )
        )
        for item in results
    )

    summary: dict[
        str,
        Any,
    ] = {
        "total": total,
        "completed": completed,
        "errors": len(errors),
        "tool_call_rate": round(
            safe_mean(tool_called_values),
            6,
        ),
        "answer_contains_score": round(
            safe_mean(answer_scores),
            6,
        ),
        "source_url_score": round(
            safe_mean(source_scores),
            6,
        ),
        "should_answer_score": round(
            safe_mean(should_answer_scores),
            6,
        ),
        "overall_score": round(
            safe_mean(overall_scores),
            6,
        ),
        "strict_pass_count": (strict_passes),
        "strict_pass_rate": round(
            (strict_passes / total if total else 0.0),
            6,
        ),
        "mean_latency_seconds": round(
            safe_mean(latencies),
            6,
        ),
        "prompt_tokens": (prompt_tokens),
        "completion_tokens": (completion_tokens),
        "total_tokens": (total_tokens),
    }

    if include_groups:
        summary["by_category"] = aggregate_by_field(
            results,
            "category",
        )

        summary["by_difficulty"] = aggregate_by_field(
            results,
            "difficulty",
        )

    return summary


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

    for result in results:
        group_name = str(
            result.get(
                field,
                "unknown",
            )
        )

        grouped[group_name].append(result)

    return {
        group_name: summarize_results(
            group_rows,
            include_groups=False,
        )
        for (
            group_name,
            group_rows,
        ) in sorted(grouped.items())
    }