"""Compare Jina Reader outputs (run2/*.json) with our own ground truth (gt_*). Pre-registered metrics: see data/tool-evidence/jina-reader/00-test-plan.md."""
import json, re, html
from html.parser import HTMLParser
def norm(s):
    URL = r"\((?:[^()\s]|\([^()]*\))+(?:\s+\"[^\"]*\")?\)"  # (url "optional title"), url may contain one level of parentheses
    s = re.sub(r"!\[[^\]]*\]" + URL, " ", s)                   # images
    s = re.sub(r"\[\[(?:\d+|[a-z]|edit|note \d+)\]\]" + URL + r"|\[\d+\]", " ", s)  # footnote / edit markers
    s = re.sub(r"\[_?\[citation needed\]" + URL + r"_?\]", " ", s)
    s = re.sub(r"\]" + URL, "]", s)                               # links -> text
    if DEHYPH: s = re.sub(r"(?<=[a-z])- ?(?=[a-z])", "", s)       # PDF line-break hyphens (reported separately)
    s = html.unescape(s).lower().replace("’", "'").replace("“", '"').replace("”", '"')
    return " ".join(re.findall(r"[a-z0-9]+", s))
DEHYPH = False
def sents(t): return [x.strip() for x in re.split(r"(?<=[.!?])\s+", t) if len(x.strip()) >= 40]
def recall(gt_sents, out):
    o = norm(out); hit = [s for s in gt_sents if norm(s) in o]
    return len(hit), len(gt_sents), [s for s in gt_sents if norm(s) not in o]
J = lambda p: json.load(open(f"run2/{p}.json"))["data"]
R = {}
# P1
gt = json.load(open("gt_p1_extract.json"))["query"]["pages"][0]["extract"]
gt = re.sub(r"\n=+ [^=]+ =+\n", "\n", gt)
for name in ("p1", "p1_nocache"):
    c = J(name)["content"]; h, n, miss = recall(sents(gt), c)
    links = {m.split("#")[0] for m in re.findall(r"\]\((https://en\.wikipedia\.org/wiki/[^)\s]+)", c) if not re.search(r"/wiki/(File|Special|Help|Wikipedia|Talk|Category|Portal|Template|Template_talk|Main_Page)[:_]?", m)}
    marks = {k: c.count(k) for k in ["Jump to content", "Main menu", "Create account", "Log in", "Donate", "Privacy policy", "Powered by MediaWiki"]}
    R[name] = {"sentence_recall": f"{h}/{n}", "pct": round(100 * h / n, 1), "missed_examples": miss[:5], "distinct_article_links": len(links), "gt_article_links": 191,
               "images_in_output": c.count("!["), "boilerplate_markers": marks, "chars": len(c), "tokens": J(name)["usage"]["tokens"]}
# images in our own article HTML
ph = json.load(open("gt_p1_links.json"))["parse"]["text"]
R["p1_gt_images"] = {"img_tags_in_article_html": ph.count("<img "), "prop_images": len(json.load(open("gt_p1_links.json"))["parse"]["images"])}
# P2
class Main(HTMLParser):
    def __init__(s): super().__init__(); s.d = 0; s.t = []; s.skip = 0
    def handle_starttag(s, tag, a):
        a = dict(a)
        if s.d: s.d += 1
        elif tag == "div" and a.get("role") == "main": s.d = 1
        if s.d and tag in ("script", "style"): s.skip += 1
    def handle_endtag(s, tag):
        if s.d:
            if tag in ("script", "style") and s.skip: s.skip -= 1
            s.d -= 1
    def handle_data(s, x):
        if s.d and not s.skip: s.t.append(x)
p = Main(); p.feed(open("gt_p2.html").read()); gt2 = " ".join(" ".join(p.t).split())
c = J("p2")["content"]; h, n, miss = recall(sents(gt2), c)
R["p2"] = {"sentence_recall": f"{h}/{n}", "pct": round(100 * h / n, 1), "missed_examples": miss[:5], "doctest_lines_out": sum(1 for l in c.splitlines() if ">>> " in l), "doctest_lines_gt": 42,
           "boilerplate_markers": {k: c.count(k) for k in ["Navigation", "Theme", "previous |", "Report a bug", "Show source", "Copyright"]}, "chars": len(c), "tokens": J("p2")["usage"]["tokens"]}
# P3
gq = [json.loads('"' + x + '"') for x in re.findall(r'"text": *"((?:[^"\\]|\\.)*)"', open("gt_p3.html").read())]
ga = [json.loads('"' + x + '"') for x in re.findall(r'"name": *"([^"]*)"', open("gt_p3.html").read())]
for name in ("p3", "p3_nocache", "p3_wait"):
    c = J(name)["content"]
    R[name] = {"quotes_found": sum(1 for q in gq if norm(q) in norm(c)), "of": len(gq), "warning": J(name).get("warning"), "tokens": J(name)["usage"]["tokens"],
               "quote_author_pairs_found": sum(1 for q, a in zip(gq, ga) if norm(q + " by " + a) in norm(c))}
R["p3_gt_authors"] = ga
# P4
ab = open("gt_p4_abs.html").read(); m = re.search(r'<meta name="citation_abstract" content="([^"]*)"', ab)
abstract = html.unescape(m.group(1)) if m else ""
c = J("p4")["content"]; h, n, miss = recall(sents(abstract), c)
DEHYPH = True; h2, _, _ = recall([re.sub(r"(?<=[a-z])-(?=[a-z])", "", x) for x in sents(abstract)], c); DEHYPH = False
R["p4"] = {"abstract_sentence_recall": f"{h}/{n}", "abstract_sentence_recall_after_removing_pdf_hyphens": f"{h2}/{n}", "hyphen_breaks_in_output": len(re.findall(r"[a-z]-[a-z]", c)), "missed": miss, "numPages": J("p4").get("numPages"), "title": J("p4")["title"], "chars": len(c), "tokens": J("p4")["usage"]["tokens"]}
json.dump(R, open("run2/metrics.json", "w"), indent=1, ensure_ascii=False); print(json.dumps(R, indent=1, ensure_ascii=False))
