"""
Comparison utilities for evaluating agent responses against expected behavior.

All functions are standalone, stdlib-only, and can be used independently or
as overrides when constructing an EvalRunner.
"""

import json
import math
import re
import string
from collections import Counter


def strip_markdown(text: str) -> str:
    """Remove markdown formatting so similarity focuses on content not style."""
    # Remove bold/italic markers
    text = re.sub(r"\*{1,3}([^*]+)\*{1,3}", r"\1", text)
    # Remove markdown headers
    text = re.sub(r"^#{1,6}\s+", "", text, flags=re.MULTILINE)
    # Remove bullet points
    text = re.sub(r"^[\s]*[-*+]\s+", "", text, flags=re.MULTILINE)
    # Remove markdown links [text](url) -> text
    text = re.sub(r"\[([^\]]+)\]\([^)]+\)", r"\1", text)
    # Collapse multiple whitespace
    text = re.sub(r"\s+", " ", text).strip()
    return text


def tokenize(text: str) -> list[str]:
    """Lowercase, strip markdown/punctuation, split into tokens."""
    text = strip_markdown(text)
    text = text.lower()
    text = text.translate(str.maketrans("", "", string.punctuation))
    return text.split()


def fuzzy_str_match(a: str, b: str, threshold: float | None = None) -> bool:
    """Check if two strings are semantically close enough to count as matching.

    Uses adaptive word-overlap thresholds (0.35-0.55) based on string length
    when no explicit threshold is provided.

    Args:
        a: First string.
        b: Second string.
        threshold: Optional explicit overlap threshold. If None, uses adaptive
            thresholds: 0.35 for <=5 words, 0.40 for <=8 words, 0.55 otherwise.

    Returns:
        True if the strings match closely enough.
    """
    na = re.sub(r"[^\w\s]", "", a.lower()).strip()
    nb = re.sub(r"[^\w\s]", "", b.lower()).strip()
    if na == nb:
        return True
    # Check if one contains the other
    if na in nb or nb in na:
        return True
    # Word-set overlap
    wa, wb = set(na.split()), set(nb.split())
    if not wa or not wb:
        return False
    overlap = len(wa & wb) / max(len(wa), len(wb))
    if threshold is None:
        min_len = min(len(wa), len(wb))
        threshold = 0.35 if min_len <= 5 else (0.40 if min_len <= 8 else 0.55)
    return overlap >= threshold


def extract_tool_args(tool_call: dict) -> dict:
    """Extract arguments from a tool call dict, handling both formats.

    Handles:
        - ``{"arguments": {...}}`` (dict form)
        - ``{"arguments_json": "..."}`` (JSON string form)

    Returns:
        The parsed arguments dict, or {} on failure.
    """
    # Prefer the parsed dict form
    if isinstance(tool_call.get("arguments"), dict):
        return tool_call["arguments"]
    # Fall back to parsing the JSON string
    raw = tool_call.get("arguments_json", "{}")
    try:
        return json.loads(raw) if isinstance(raw, str) else (raw or {})
    except (json.JSONDecodeError, TypeError):
        return {}


def compare_tool_args(
    expected: dict, actual: dict
) -> tuple[str, str | None]:
    """Compare expected vs actual tool arguments.

    Returns:
        A tuple of (match_status, divergence_notes) where match_status is one
        of "exact", "partial", or "mismatch", and divergence_notes is a
        human-readable summary of differences (or None if exact).
    """
    exp_args = extract_tool_args(expected)
    act_args = extract_tool_args(actual)

    if not exp_args and not act_args:
        return "exact", None

    all_match = True
    any_match = False
    for key, exp_val in exp_args.items():
        act_val = act_args.get(key)
        if act_val is None:
            all_match = False
            continue
        if isinstance(exp_val, str) and isinstance(act_val, str):
            if fuzzy_str_match(exp_val, act_val):
                any_match = True
            else:
                all_match = False
        else:
            if exp_val == act_val:
                any_match = True
            else:
                all_match = False

    if all_match and exp_args:
        status = "exact"
    elif any_match:
        status = "partial"
    elif exp_args:
        status = "mismatch"
    else:
        status = "exact"

    notes = _arg_diff_notes(exp_args, act_args) if status != "exact" else None
    return status, notes


def text_similarity(text_a: str, text_b: str) -> float:
    """Compute similarity between two text strings.

    Uses cosine similarity on word frequency vectors, plus bonuses for
    matching entities (order IDs, prices, dates) and domain concepts
    (refund, shipped, in stock, etc.).

    Returns:
        A float between 0.0 and 1.0.
    """
    tokens_a = tokenize(text_a)
    tokens_b = tokenize(text_b)

    if not tokens_a or not tokens_b:
        return 0.0

    counter_a = Counter(tokens_a)
    counter_b = Counter(tokens_b)
    all_words = set(counter_a.keys()) | set(counter_b.keys())

    dot = sum(counter_a.get(w, 0) * counter_b.get(w, 0) for w in all_words)
    mag_a = math.sqrt(sum(v * v for v in counter_a.values()))
    mag_b = math.sqrt(sum(v * v for v in counter_b.values()))
    cosine = dot / (mag_a * mag_b) if mag_a > 0 and mag_b > 0 else 0.0

    # Boost 1: Entity overlap (order IDs, refund IDs, prices, dates, tracking)
    entity_pattern = re.compile(
        r"ORD-\d+|REF-\d+|\$[\d.]+|\d{4}-\d{2}-\d{2}|track\S+\.\w+/\S+",
        re.IGNORECASE,
    )
    entities_a = set(entity_pattern.findall(text_a.upper()))
    entities_b = set(entity_pattern.findall(text_b.upper()))
    if entities_a and entities_b:
        entity_overlap = len(entities_a & entities_b) / max(
            len(entities_a), len(entities_b)
        )
        cosine = min(1.0, cosine + entity_overlap * 0.20)

    # Boost 2: Domain concept overlap
    concepts = [
        {"refund", "refunded", "credited"},
        {"shipped", "shipping", "transit", "delivered", "delivery"},
        {"tracking", "track"},
        {"stock", "available", "availability", "inventory"},
        {"processing", "processed"},
        {"pickup", "store"},
        {"manual", "review", "escalat"},
        {"damaged", "defective", "cracked", "broken"},
    ]
    la, lb = text_a.lower(), text_b.lower()
    concept_matches = 0
    concept_total = 0
    for concept_set in concepts:
        a_has = any(c in la for c in concept_set)
        b_has = any(c in lb for c in concept_set)
        if a_has or b_has:
            concept_total += 1
            if a_has and b_has:
                concept_matches += 1
    if concept_total > 0:
        concept_overlap = concept_matches / concept_total
        cosine = min(1.0, cosine + concept_overlap * 0.10)

    return round(cosine, 2)


def compare_args_structural(
    expected: dict, actual: dict
) -> tuple[str, dict]:
    """Compute a structural argument diff for the backend grader.

    Produces an ``argument_comparison`` dict with ``matching``, ``different``,
    ``missing``, and ``extra`` buckets — the exact format the backend's
    ``regrade_tool_args()`` expects.  No fuzzy matching or scoring is done;
    this is a literal key-by-key comparison.

    Args:
        expected: Expected tool call dict (with ``arguments`` or ``arguments_json``).
        actual: Actual tool call dict (same format).

    Returns:
        A tuple of ``(match_status, argument_comparison)`` where match_status
        is ``"exact"`` if all args match literally, ``"mismatch"`` if tool
        names differ or one side is NOT_CALLED/NONE_EXPECTED, or ``"partial"``
        if tool names match but arguments differ.
    """
    exp_name = expected.get("name") or expected.get("tool_name", "")
    act_name = actual.get("name") or actual.get("tool_name", "")

    empty_comp = {"matching": {}, "different": {}, "missing": {}, "extra": {}}

    # Tool not called or unexpected
    if act_name == "NOT_CALLED":
        return "mismatch", empty_comp
    if exp_name == "NONE_EXPECTED":
        return "mismatch", empty_comp
    if exp_name != act_name:
        return "mismatch", empty_comp

    exp_args = extract_tool_args(expected)
    act_args = extract_tool_args(actual)

    arg_comparison: dict[str, dict] = {
        "matching": {},
        "different": {},
        "missing": {},
        "extra": {},
    }

    all_keys = set(exp_args.keys()) | set(act_args.keys())

    for key in all_keys:
        exp_val = exp_args.get(key)
        act_val = act_args.get(key)
        if key in exp_args and key in act_args:
            if exp_val == act_val:
                arg_comparison["matching"][key] = {"expected": exp_val, "actual": act_val}
            else:
                arg_comparison["different"][key] = {"expected": exp_val, "actual": act_val}
        elif key in exp_args:
            arg_comparison["missing"][key] = {"expected": exp_val}
        else:
            arg_comparison["extra"][key] = {"actual": act_val}

    if not arg_comparison["different"] and not arg_comparison["missing"] and not arg_comparison["extra"]:
        return "exact", arg_comparison

    return "partial", arg_comparison


def _arg_diff_notes(exp_args: dict, act_args: dict) -> str | None:
    """Generate human-readable notes about argument differences."""
    diffs = []
    for key in set(list(exp_args.keys()) + list(act_args.keys())):
        ev = exp_args.get(key)
        av = act_args.get(key)
        if ev is None:
            diffs.append(f"extra arg '{key}'={av}")
        elif av is None:
            diffs.append(f"missing arg '{key}'")
        elif isinstance(ev, str) and isinstance(av, str):
            if not fuzzy_str_match(ev, av):
                diffs.append(f"'{key}': expected='{ev}' actual='{av}'")
        elif ev != av:
            diffs.append(f"'{key}': expected={ev} actual={av}")
    return "; ".join(diffs) if diffs else None
