"""Assertion-completion scoring, identical to RepoPeftBench's evaluation/metrics.py
(exact match and edit similarity; CodeBLEU omitted)."""

from difflib import SequenceMatcher

FIM_TOKENS = ("<|fim_prefix|>", "<|fim_suffix|>", "<|fim_middle|>")


def strip_fim_tokens(s: str) -> str:
    for tok in FIM_TOKENS:
        s = s.replace(tok, "")
    return s.strip()


def strip_comments(s: str) -> str:
    return s.split("#")[0].strip()


def normalize_for_match(s: str) -> str:
    s = s.strip().rstrip(":, \t")
    s = s.replace(", ", ",").replace(" ,", ",")
    return " ".join(s.split())


def exact_match(pred: str, ref: str) -> bool:
    return normalize_for_match(pred) == normalize_for_match(ref)


def edit_similarity(pred: str, ref: str) -> float:
    return SequenceMatcher(None, pred, ref).ratio()


def postprocess_prediction(pred: str, target: str) -> str:
    pred = strip_fim_tokens(pred)
    if "\n" not in target and "\n" in pred:
        pred = pred.split("\n")[0]
    pred = strip_comments(pred)
    ref = strip_comments(target).strip()
    if ref:
        if "," not in ref and "," in pred:
            pred = pred.split(",")[0].strip()
        if len(ref.split()) == 1 and len(pred.strip().split()) > 1:
            pred = pred.strip().split()[0]
    return pred.strip()


def compute_metrics(pred: str, ref: str) -> dict:
    pred_clean = postprocess_prediction(pred, ref)
    ref_clean = strip_comments(ref)
    return {
        "pred_clean": pred_clean,
        "exact_match": exact_match(pred_clean, ref_clean),
        "edit_similarity": edit_similarity(pred_clean, ref_clean),
    }