"""
Flag Stage 1 paragraphs needing Stage 2 re-evaluation due to codebook v2.5->v3.5 drift.

Two categories of flags:
1. Materiality assessment language in N/O or RMP paragraphs — backward-looking
   conclusions or SEC-qualified forward-looking statements that constitute a
   materiality assessment (should be Strategy Integration under v3.5 codebook).
2. SPAC/shell company paragraphs coded as substantive categories — should be None/Other.

Materiality rule (tightened after 6 rounds of prompt iteration):
  IS assessment: "have not materially affected/impacted", "not materially affected",
    "reasonably likely to materially affect/impact",
    "have not experienced any material cybersecurity" (unless cross-reference context).
  NOT assessment: "could/may ... material adverse effect" (boilerplate speculation),
    "material" as adjective ("material risks"), cross-references ("see Item 1A"),
    consequence clauses at end of RMP descriptions.

Usage: uv run scripts/flag-stage1-corrections.py
"""

import json
import re
from collections import Counter, defaultdict
from pathlib import Path

DATA_DIR = Path(__file__).resolve().parent.parent / "data"
STAGE1_PATH = DATA_DIR / "annotations" / "stage1.patched.jsonl"
PARAGRAPHS_PATH = DATA_DIR / "paragraphs" / "paragraphs-clean.patched.jsonl"
HOLDOUT_PATH = DATA_DIR / "gold" / "human-labels-raw.jsonl"
OUTPUT_PATH = DATA_DIR / "annotations" / "stage1-corrections.jsonl"

# Category abbreviation mapping for stage1Labels output
CATEGORY_ABBREV = {
    "None/Other": "N/O",
    "Board Governance": "BG",
    "Management Role": "MR",
    "Risk Management Process": "RMP",
    "Third-Party Risk": "TPR",
    "Incident Disclosure": "ID",
}

# --- Materiality patterns (strict assessment-only) ---

# Positive patterns: backward-looking conclusions and SEC-qualified forward-looking
MATERIALITY_ASSESSMENT_RE = re.compile(
    r"(?:"
    # Backward-looking conclusions
    r"have\s+not\s+materially\s+affected"
    r"|has\s+not\s+materially\s+affected"
    r"|not\s+materially\s+affected"
    r"|have\s+not\s+materially\s+impacted"
    # SEC-qualified forward-looking
    r"|reasonably\s+likely\s+to\s+materially\s+affect"
    r"|reasonably\s+likely\s+to\s+materially\s+impact"
    # Negative assertions about incidents
    r"|have\s+not\s+experienced\s+any\s+material\s+cybersecurity"
    r")",
    re.IGNORECASE,
)

# Negative filter: cross-reference context near the match (within 200 chars after)
CROSS_REF_RE = re.compile(
    r"(?:see\s+Item|see\s+Part|see\s+our\s+risk\s+factors|refer\s+to)",
    re.IGNORECASE,
)

# Negative filter: speculative/boilerplate "could/may + material adverse" patterns
SPECULATIVE_RE = re.compile(
    r"(?:could|may|might|can)\s+(?:\w+\s+){0,3}material\s+adverse\s+effect",
    re.IGNORECASE,
)


def has_materiality_language(text: str) -> str | None:
    """Return the matched materiality assessment pattern string, or None if no match."""
    match = MATERIALITY_ASSESSMENT_RE.search(text)
    if match is None:
        return None

    matched_text = match.group(0)
    match_start = match.start()
    match_end = match.end()

    # Check cross-reference context: look within 200 chars after the match
    post_context = text[match_end : match_end + 200]
    if CROSS_REF_RE.search(post_context):
        return None

    # Also check cross-reference before the match (within 100 chars)
    pre_context = text[max(0, match_start - 100) : match_start]
    if CROSS_REF_RE.search(pre_context):
        return None

    # Extract a snippet around the match for context
    snippet_start = max(0, match_start - 30)
    snippet_end = min(len(text), match_end + 30)
    snippet = text[snippet_start:snippet_end].strip()

    return snippet


# --- SPAC patterns ---

SPAC_PHRASES = [
    "special purpose acquisition",
    "blank check",
    "no business operations",
    "shell company",
    "have not adopted any cybersecurity",
    "no operations",
]


def has_spac_language(text: str) -> str | None:
    """Return matched SPAC indicator string, or None."""
    text_lower = text.lower()
    for phrase in SPAC_PHRASES:
        if phrase in text_lower:
            idx = text_lower.index(phrase)
            start = max(0, idx - 20)
            end = min(len(text), idx + len(phrase) + 20)
            return text[start:end].strip()
    return None


def main() -> None:
    # Load holdout IDs
    print("Loading holdout IDs...")
    holdout_ids: set[str] = set()
    with open(HOLDOUT_PATH) as f:
        for line in f:
            rec = json.loads(line)
            holdout_ids.add(rec["paragraphId"])
    print(f"  Holdout paragraphs: {len(holdout_ids)}")

    # Load paragraph texts
    print("Loading paragraph texts...")
    para_texts: dict[str, str] = {}
    with open(PARAGRAPHS_PATH) as f:
        for line in f:
            rec = json.loads(line)
            para_texts[rec["id"]] = rec["text"]
    print(f"  Paragraphs loaded: {len(para_texts)}")

    # Load old corrections for comparison
    old_materiality_pids: set[str] = set()
    old_spac_pids: set[str] = set()
    if OUTPUT_PATH.exists():
        with open(OUTPUT_PATH) as f:
            for line in f:
                rec = json.loads(line)
                if rec["reason"] == "materiality_language":
                    old_materiality_pids.add(rec["paragraphId"])
                elif rec["reason"] == "spac":
                    old_spac_pids.add(rec["paragraphId"])

    # Load Stage 1 annotations and group by paragraphId
    print("Loading Stage 1 annotations...")
    annotations: dict[str, list[str]] = defaultdict(list)
    with open(STAGE1_PATH) as f:
        for line in f:
            rec = json.loads(line)
            pid = rec["paragraphId"]
            if pid in holdout_ids:
                continue
            cat = rec["label"]["content_category"]
            annotations[pid].append(cat)

    total_paragraphs = len(annotations)
    print(f"  Stage 1 paragraphs (excluding holdout): {total_paragraphs}")

    # Process each paragraph
    flagged: list[dict] = []
    materiality_flagged: list[dict] = []
    spac_flagged: list[dict] = []

    # Track paragraphs that HAVE assessment language but pass through new filters
    # (for showing "newly excluded" examples)
    newly_excluded: list[dict] = []

    for pid, labels in annotations.items():
        text = para_texts.get(pid)
        if text is None:
            continue

        label_abbrevs = [CATEGORY_ABBREV.get(l, l) for l in labels]
        no_count = sum(1 for l in labels if l == "None/Other")
        total = len(labels)

        # --- Check 1: Materiality assessment in N/O paragraphs ---
        if no_count > total / 2:  # majority or unanimous N/O
            matched = has_materiality_language(text)
            if matched:
                is_unanimous = no_count == total
                consensus = "unanimous" if is_unanimous else "majority"
                record = {
                    "paragraphId": pid,
                    "reason": "materiality_language",
                    "originalConsensus": consensus,
                    "originalCategory": "None/Other",
                    "matchedPattern": matched,
                    "stage1Labels": label_abbrevs,
                }
                flagged.append(record)
                materiality_flagged.append(record)
            elif pid in old_materiality_pids:
                # Was flagged before, now excluded — collect for comparison
                # Find what the old broad matcher would have caught
                text_lower = text.lower()
                broad_match = None
                for phrase in [
                    "material adverse", "materially affect", "material impact",
                    "material cybersecurity", "material effect", "not materially",
                    "materially impacted", "materially affected",
                ]:
                    if phrase in text_lower:
                        idx = text_lower.index(phrase)
                        s = max(0, idx - 30)
                        e = min(len(text), idx + len(phrase) + 30)
                        broad_match = text[s:e].strip()
                        break
                if broad_match is None:
                    broad_match = "(proximity match)"
                newly_excluded.append({
                    "paragraphId": pid,
                    "reason": "excluded_materiality",
                    "oldMatch": broad_match,
                    "stage1Labels": label_abbrevs,
                })

        # --- Check 2: SPAC paragraphs coded as non-N/O ---
        if no_count <= total / 2:
            matched = has_spac_language(text)
            if matched:
                cat_counts = Counter(labels)
                majority_cat = cat_counts.most_common(1)[0][0]
                consensus = "unanimous" if cat_counts[majority_cat] == total else "majority"
                record = {
                    "paragraphId": pid,
                    "reason": "spac",
                    "originalConsensus": consensus,
                    "originalCategory": majority_cat,
                    "matchedPattern": matched,
                    "stage1Labels": label_abbrevs,
                }
                flagged.append(record)
                spac_flagged.append(record)

    # Write output
    print(f"\nWriting {len(flagged)} flagged paragraphs to {OUTPUT_PATH}...")
    with open(OUTPUT_PATH, "w") as f:
        for rec in flagged:
            f.write(json.dumps(rec) + "\n")

    # --- Print summary ---
    print("\n" + "=" * 70)
    print("STAGE 1 CORRECTION FLAGS — SUMMARY")
    print("=" * 70)
    print(f"Total Stage 1 paragraphs (excluding holdout): {total_paragraphs:,}")
    print()
    print("  Comparison (old broad rule -> new strict rule):")
    print(f"    Materiality flags:  {len(old_materiality_pids):>5}  ->  {len(materiality_flagged):>5}  (delta: {len(materiality_flagged) - len(old_materiality_pids):+d})")
    print(f"    SPAC flags:         {len(old_spac_pids):>5}  ->  {len(spac_flagged):>5}  (delta: {len(spac_flagged) - len(old_spac_pids):+d})")
    old_total = len(old_materiality_pids) + len(old_spac_pids)
    print(f"    Total flags:        {old_total:>5}  ->  {len(flagged):>5}  (delta: {len(flagged) - old_total:+d})")
    print()

    new_materiality_pids = {r["paragraphId"] for r in materiality_flagged}
    added = new_materiality_pids - old_materiality_pids
    removed = old_materiality_pids - new_materiality_pids
    kept = new_materiality_pids & old_materiality_pids
    print(f"  Materiality breakdown:")
    print(f"    Kept from old:      {len(kept):>5}")
    print(f"    Newly flagged:      {len(added):>5}")
    print(f"    Excluded (dropped): {len(removed):>5}")

    # Show examples
    def show_examples(title: str, records: list[dict], n: int = 5, text_key: str = "matchedPattern") -> None:
        print(f"\n--- {title} (showing {min(n, len(records))} of {len(records)}) ---")
        for rec in records[:n]:
            pid = rec["paragraphId"]
            text = para_texts.get(pid, "")
            snippet = text[:150] + "..." if len(text) > 150 else text
            print(f"  {pid[:16]}...")
            print(f"    Labels: {rec['stage1Labels']}")
            print(f"    Match:  {rec.get(text_key, rec.get('oldMatch', ''))}")
            print(f"    Text:   {snippet}")
            print()

    show_examples(
        "Newly flagged materiality assessments (assessment patterns)",
        [r for r in materiality_flagged if r["paragraphId"] in added] or materiality_flagged,
    )
    show_examples(
        "Previously flagged, NOW EXCLUDED (boilerplate/speculation)",
        newly_excluded,
        text_key="oldMatch",
    )
    show_examples(
        "SPAC/shell coded as non-N/O", spac_flagged
    )


if __name__ == "__main__":
    main()
