from __future__ import annotations

import argparse
import csv
import math
import re
import statistics
from collections import Counter
from pathlib import Path

import torch
from tokenizers import Tokenizer

from build_owt_t5_stream_clean_pack_cache import (
    CODE_SYMBOLS,
    LIST_LINE_RE,
    PUNCT_SYMBOLS,
    SENTENCE_RE,
    WORD_RE,
    char_run_info,
    max_run_info,
    repetitive_pattern_info,
    repeated_meaningful_ngram_info,
    text_quality_reasons,
    top_token_info,
)


STOPWORDS = {
    "a",
    "about",
    "after",
    "all",
    "also",
    "an",
    "and",
    "are",
    "as",
    "at",
    "be",
    "because",
    "been",
    "but",
    "by",
    "can",
    "could",
    "for",
    "from",
    "had",
    "has",
    "have",
    "he",
    "her",
    "his",
    "i",
    "if",
    "in",
    "into",
    "is",
    "it",
    "its",
    "more",
    "not",
    "of",
    "on",
    "one",
    "or",
    "our",
    "she",
    "so",
    "some",
    "than",
    "that",
    "the",
    "their",
    "them",
    "there",
    "these",
    "they",
    "this",
    "to",
    "was",
    "we",
    "were",
    "what",
    "when",
    "which",
    "who",
    "will",
    "with",
    "would",
    "you",
    "your",
}


def q(values: list[float], p: float) -> float:
    if not values:
        return float("nan")
    values = sorted(values)
    return values[min(len(values) - 1, max(0, int(round(p * (len(values) - 1)))))]


def short(text: str, n: int = 180) -> str:
    return text.replace("\n", "\\n").replace("\t", "\\t")[:n]


def token_text(tok: Tokenizer, idx: int) -> str:
    return tok.decode([int(idx)], skip_special_tokens=False)


def char_frac(text: str, pred) -> float:
    visible = [ch for ch in text if not ch.isspace()]
    if not visible:
        return 0.0
    return sum(1 for ch in visible if pred(ch)) / len(visible)


def row_metrics(row: list[int], tok: Tokenizer, eos_id: int, unk_id: int) -> dict[str, object]:
    # Keep all tokens for positional/special stats; use payload for content stats.
    payload = row[:]
    text = tok.decode(payload, skip_special_tokens=False)
    words = WORD_RE.findall(text)
    low_words = [w.lower().strip("'") for w in words]
    stop_frac = sum(1 for w in low_words if w in STOPWORDS) / max(1, len(low_words))
    line_items = [line.strip() for line in text.splitlines() if line.strip()]
    line_count = len(line_items)
    short_line_frac = sum(1 for line in line_items if len(line) <= 48) / max(1, line_count)
    list_line_frac = sum(1 for line in line_items if LIST_LINE_RE.search(line)) / max(1, line_count)
    table_line_frac = sum(1 for line in line_items if "|" in line or "\t" in line or re.search(r"\s{4,}", line)) / max(
        1, line_count
    )
    top_count, top_frac, top_token = top_token_info(payload, tok)
    run_count, run_token = max_run_info(payload, tok)
    bigram_count, bigram_text = repeated_meaningful_ngram_info(payload, 2, tok, threshold=8)
    trigram_count, trigram_text = repeated_meaningful_ngram_info(payload, 3, tok, threshold=8)
    pattern_repeats, pattern_width, pattern_text = repetitive_pattern_info(payload, tok)
    comma_run, comma_text = char_run_info(text, r",+")
    apostrophe_run, apostrophe_text = char_run_info(text, r"'+")
    quote_run, quote_text = char_run_info(text, r"\"+")
    dash_run, dash_text = char_run_info(text, r"-+")
    dot_run, dot_text = char_run_info(text, r"\.+")
    punct_run, punct_text = char_run_info(text, r"[-_=*~.]{2,}|[!?]{2,}")
    sentence_count = len(SENTENCE_RE.findall(text))
    unique = len(set(payload))
    eos_count = payload.count(eos_id)
    unk_count = payload.count(unk_id)
    alpha_frac = char_frac(text, str.isalpha)
    digit_frac = char_frac(text, str.isdigit)
    punct_frac = char_frac(text, lambda ch: ch in PUNCT_SYMBOLS)
    code_symbol_frac = char_frac(text, lambda ch: ch in CODE_SYMBOLS)
    non_ascii_frac = char_frac(text, lambda ch: ord(ch) > 127)

    strict_reasons = text_quality_reasons(
        text,
        max_html_entities=0,
        reject_code_like=True,
        reject_list_like=True,
        prose_like=True,
        max_comma_run=1,
        max_apostrophe_run=1,
        max_quote_run=1,
        max_punct_run=7,
        max_punct_frac=0.20,
        max_code_symbol_frac=0.08,
        min_alpha_frac=0.62,
        min_words=64,
        min_sentences=2,
        max_url_count=2,
    )

    flags: list[str] = []
    if unk_count:
        flags.append(f"unk={unk_count}")
    if unique <= 80:
        flags.append(f"low_unique={unique}")
    if top_frac >= 0.06:
        flags.append(f"top_token={top_frac:.3f}:{top_token}")
    if run_count >= 6:
        flags.append(f"run={run_count}:{run_token}")
    if bigram_count >= 8:
        flags.append(f"bigram={bigram_count}:{bigram_text}")
    if trigram_count >= 8:
        flags.append(f"trigram={trigram_count}:{trigram_text}")
    if pattern_repeats > 3:
        flags.append(f"pattern={pattern_repeats}x{pattern_width}:{pattern_text}")
    if punct_frac >= 0.16:
        flags.append(f"punct_frac={punct_frac:.3f}")
    if code_symbol_frac >= 0.055:
        flags.append(f"code_symbol_frac={code_symbol_frac:.3f}")
    if stop_frac <= 0.18 and len(words) >= 80:
        flags.append(f"low_stop={stop_frac:.3f}")
    if non_ascii_frac >= 0.025:
        flags.append(f"non_ascii={non_ascii_frac:.3f}")
    if line_count >= 8 and (short_line_frac >= 0.5 or list_line_frac >= 0.25 or table_line_frac >= 0.25):
        flags.append(f"list_table_lines={line_count}:{short_line_frac:.2f}/{list_line_frac:.2f}/{table_line_frac:.2f}")
    if comma_run > 1:
        flags.append(f"comma_run={comma_run}:{comma_text}")
    if apostrophe_run > 1:
        flags.append(f"apostrophe_run={apostrophe_run}:{apostrophe_text}")
    if quote_run > 1:
        flags.append(f"quote_run={quote_run}:{quote_text}")
    if dash_run > 7 or dot_run > 7 or punct_run > 7:
        flags.append(f"punct_run={max(dash_run, dot_run, punct_run)}:{punct_text or dash_text or dot_text}")
    for reason in strict_reasons:
        if reason not in flags:
            flags.append(reason)

    # The score is only for sorting suspicious examples, not a training metric.
    score = 0.0
    score += max(0.0, top_frac - 0.035) * 16.0
    score += max(0.0, punct_frac - 0.10) * 4.0
    score += max(0.0, code_symbol_frac - 0.035) * 5.0
    score += max(0.0, 0.22 - stop_frac) * 2.5 if len(words) >= 80 else 0.0
    score += max(0.0, non_ascii_frac - 0.01) * 3.0
    score += min(0.4, max(0, bigram_count - 4) / 20)
    score += min(0.4, max(0, trigram_count - 4) / 20)
    score += min(0.3, max(0, run_count - 3) / 20)
    score += min(0.4, len(flags) / 20)

    return {
        "len": len(payload),
        "first_id": payload[0],
        "first_tok": short(token_text(tok, payload[0]), 40),
        "last_id": payload[-1],
        "last_tok": short(token_text(tok, payload[-1]), 40),
        "eos_count": eos_count,
        "internal_eos_count": max(0, eos_count - int(payload[-1] == eos_id)),
        "unk_count": unk_count,
        "unique": unique,
        "top_count": top_count,
        "top_frac": top_frac,
        "top_token": top_token,
        "max_run": run_count,
        "max_run_token": run_token,
        "bigram_count": bigram_count,
        "bigram_text": bigram_text,
        "trigram_count": trigram_count,
        "trigram_text": trigram_text,
        "pattern_repeats": pattern_repeats,
        "pattern_width": pattern_width,
        "pattern_text": pattern_text,
        "char_len": len(text),
        "word_count": len(words),
        "sentence_count": sentence_count,
        "alpha_frac": alpha_frac,
        "digit_frac": digit_frac,
        "punct_frac": punct_frac,
        "code_symbol_frac": code_symbol_frac,
        "non_ascii_frac": non_ascii_frac,
        "stop_frac": stop_frac,
        "line_count": line_count,
        "short_line_frac": short_line_frac,
        "list_line_frac": list_line_frac,
        "table_line_frac": table_line_frac,
        "comma_run": comma_run,
        "apostrophe_run": apostrophe_run,
        "quote_run": quote_run,
        "dash_run": dash_run,
        "dot_run": dot_run,
        "punct_run": punct_run,
        "flags": "|".join(flags),
        "flag_count": len(flags),
        "score": score,
        "preview": short(text, 240),
        "text": text,
    }


def summarize(rows: list[dict[str, object]], key: str) -> str:
    vals = [float(r[key]) for r in rows]
    return (
        f"{key}: mean={statistics.mean(vals):.4g} p50={q(vals, 0.50):.4g} "
        f"p90={q(vals, 0.90):.4g} p99={q(vals, 0.99):.4g} max={max(vals):.4g}"
    )


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--cache_path", required=True)
    p.add_argument("--tokenizer_path", default="/e2e-data/evad-tech-vla/wanghan58/models/hf/t5-small/tokenizer.json")
    p.add_argument("--out_dir", required=True)
    p.add_argument("--max_rows", type=int, default=0)
    p.add_argument("--worst", type=int, default=80)
    args = p.parse_args()

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    tok = Tokenizer.from_file(args.tokenizer_path)
    cache = torch.load(args.cache_path, map_location="cpu")
    ids = cache["ids"]
    eos_id = int(cache.get("eos_id", tok.token_to_id("</s>")))
    unk_id = int(cache.get("unk_id", tok.token_to_id("<unk>")))
    total = int(ids.shape[0])
    n = min(total, args.max_rows) if args.max_rows > 0 else total

    rows: list[dict[str, object]] = []
    token_counts: Counter[int] = Counter()
    first_counts: Counter[int] = Counter()
    flag_counts: Counter[str] = Counter()
    for idx in range(n):
        row = [int(x) for x in ids[idx].tolist()]
        token_counts.update(row)
        first_counts[row[0]] += 1
        m = row_metrics(row, tok, eos_id, unk_id)
        m["idx"] = idx
        rows.append(m)
        for flag in str(m["flags"]).split("|"):
            if flag:
                flag_counts[flag.split("=", 1)[0].split(">", 1)[0].split("<", 1)[0]] += 1

    metric_fields = [
        "idx",
        "score",
        "flag_count",
        "flags",
        "len",
        "first_tok",
        "last_tok",
        "eos_count",
        "internal_eos_count",
        "unk_count",
        "unique",
        "top_frac",
        "top_token",
        "max_run",
        "max_run_token",
        "bigram_count",
        "bigram_text",
        "trigram_count",
        "trigram_text",
        "pattern_repeats",
        "pattern_width",
        "pattern_text",
        "word_count",
        "sentence_count",
        "alpha_frac",
        "digit_frac",
        "punct_frac",
        "code_symbol_frac",
        "non_ascii_frac",
        "stop_frac",
        "line_count",
        "short_line_frac",
        "list_line_frac",
        "table_line_frac",
        "comma_run",
        "apostrophe_run",
        "quote_run",
        "dash_run",
        "dot_run",
        "punct_run",
        "preview",
    ]
    with (out_dir / "row_metrics.tsv").open("w", encoding="utf-8", newline="") as f:
        w = csv.DictWriter(f, fieldnames=metric_fields, delimiter="\t", extrasaction="ignore")
        w.writeheader()
        for r in rows:
            w.writerow(r)

    worst_rows = sorted(rows, key=lambda r: float(r["score"]), reverse=True)[: args.worst]
    with (out_dir / "worst_rows.txt").open("w", encoding="utf-8") as f:
        for r in worst_rows:
            f.write(
                f"===== idx={r['idx']} score={float(r['score']):.4f} flags={r['flags']} "
                f"top={r['top_frac']:.4f}:{r['top_token']} stop={r['stop_frac']:.4f} "
                f"punct={r['punct_frac']:.4f} unique={r['unique']} eos_internal={r['internal_eos_count']} =====\n"
            )
            f.write(str(r["text"]).replace("\r", "")[:6000])
            f.write("\n\n")

    summary_lines: list[str] = []
    summary_lines.append(f"cache={args.cache_path}")
    summary_lines.append(f"shape={tuple(ids.shape)} audited_rows={n} source={cache.get('source')}")
    summary_lines.append("")
    for key in [
        "score",
        "flag_count",
        "unique",
        "top_frac",
        "max_run",
        "bigram_count",
        "trigram_count",
        "word_count",
        "sentence_count",
        "alpha_frac",
        "punct_frac",
        "code_symbol_frac",
        "non_ascii_frac",
        "stop_frac",
        "internal_eos_count",
        "line_count",
        "short_line_frac",
    ]:
        summary_lines.append(summarize(rows, key))
    summary_lines.append("")
    summary_lines.append("flag_counts:")
    for flag, count in flag_counts.most_common(40):
        summary_lines.append(f"  {flag}: {count}/{n} ({count / n:.2%})")
    summary_lines.append("")
    summary_lines.append("top_tokens_all:")
    denom = sum(token_counts.values())
    for tid, count in token_counts.most_common(60):
        summary_lines.append(f"  {count / denom:.5f}\t{count}\t{tid}\t{short(token_text(tok, tid), 80)}")
    summary_lines.append("")
    summary_lines.append("first_tokens:")
    for tid, count in first_counts.most_common(40):
        summary_lines.append(f"  {count / n:.5f}\t{count}\t{tid}\t{short(token_text(tok, tid), 80)}")
    summary_lines.append("")
    summary_lines.append("worst_indices:")
    for r in worst_rows[:20]:
        summary_lines.append(f"  idx={r['idx']} score={float(r['score']):.4f} flags={r['flags']} preview={r['preview']}")

    (out_dir / "summary.txt").write_text("\n".join(summary_lines) + "\n", encoding="utf-8")
    print("\n".join(summary_lines[:80]))


if __name__ == "__main__":
    main()