from __future__ import annotations

import argparse
import json
import re
import unicodedata
from difflib import SequenceMatcher
from pathlib import Path


STOPWORDS = {"a", "o", "as", "os", "de", "da", "do", "das", "dos", "e", "em", "um", "uma", "para", "por", "que", "é", "são", "como", "qual", "quais", "quando", "não", "com"}


def tokens(text: str) -> set[str]:
    normalized = unicodedata.normalize("NFKD", text.lower()).encode("ascii", "ignore").decode("ascii")
    return {token for token in re.findall(r"[a-z0-9]+", normalized) if token not in {unicodedata.normalize("NFKD", w).encode("ascii", "ignore").decode("ascii") for w in STOPWORDS} and len(token) > 2}


def score_pair(query: str, document: str) -> dict:
    q = tokens(query)
    d = tokens(document)
    exact = len(q & d) / max(1, len(q))
    fuzzy_hits = 0
    for query_token in q - d:
        if any(SequenceMatcher(None, query_token, doc_token).ratio() >= 0.72 for doc_token in d):
            fuzzy_hits += 1
    overlap = min(1.0, exact + 0.5 * fuzzy_hits / max(1, len(q)))
    numbers_q = set(re.findall(r"\d+", query))
    numbers_d = set(re.findall(r"\d+", document))
    number_bonus = 0.15 if numbers_q and numbers_q & numbers_d else 0.0
    score = min(1.0, overlap + number_bonus)
    label = 2 if score >= 0.42 else 1 if score >= 0.16 else 0
    return {"score": round(score, 6), "predicted_label": label, "overlap": round(overlap, 6), "number_bonus": number_bonus, "human_review_required": True}


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("record_json", type=Path)
    args = parser.parse_args()
    row = json.loads(args.record_json.read_text(encoding="utf-8"))
    print(json.dumps(score_pair(row["query"], row["document"]), ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()