"""Lecture JSONL redaction: instructor-only policy.

Drops all turns where speaker == "Anonymous Student". Redacts known student
first/last names from instructor turns and problem_outcome records. Whitelist
phrases preserved verbatim. Defaults loaded from config/redaction.yaml.
Atomic write via temp file + os.replace. Emits sibling redaction-report.txt.
"""
from __future__ import annotations
import json
import os
import re
from dataclasses import dataclass, field
from pathlib import Path

import yaml


REDACTION_POLICY_VERSION = "1.0"
_CONFIG_PATH = Path(__file__).parent.parent.parent / "config" / "redaction.yaml"


@dataclass
class RedactionConfig:
    first_names: set = field(default_factory=set)
    last_names: set = field(default_factory=set)
    whitelist_phrases: list = field(default_factory=list)
    replacement_token: str = "[STUDENT]"
    last_name_replacement_token: str = "[REDACTED-LNAME]"
    policy_version: str = REDACTION_POLICY_VERSION


@dataclass
class RedactionStats:
    student_turns_dropped: int
    instructor_turns_kept: int
    session_headers_kept: int
    problem_outcomes_kept: int
    name_redactions_in_instructor: int
    name_redactions_in_problem_outcomes: int


def default_config() -> RedactionConfig:
    raw = yaml.safe_load(_CONFIG_PATH.read_text(encoding="utf-8"))
    first = set(raw.get("verified_first_names", [])) | set(raw.get("defensive_first_names", []))
    return RedactionConfig(
        first_names=first,
        last_names=set(raw.get("last_names", [])),
        whitelist_phrases=list(raw.get("whitelist_phrases", [])),
        replacement_token=raw.get("replacement_token", "[STUDENT]"),
        last_name_replacement_token=raw.get("last_name_replacement_token", "[REDACTED-LNAME]"),
        policy_version=raw.get("policy_version", REDACTION_POLICY_VERSION),
    )


def _build_regexes(config):
    whitelist_rx = re.compile("|".join(re.escape(p) for p in config.whitelist_phrases)) if config.whitelist_phrases else re.compile(r"$^")
    name_rx = re.compile(
        r"\b(" + "|".join(re.escape(n) for n in sorted(config.first_names, key=len, reverse=True)) + r")\b",
        flags=re.IGNORECASE,
    ) if config.first_names else re.compile(r"$^")
    last_rx = re.compile(
        r"\b(" + "|".join(re.escape(n) for n in config.last_names) + r")\b",
        flags=re.IGNORECASE,
    ) if config.last_names else re.compile(r"$^")
    return whitelist_rx, name_rx, last_rx


def _redact_text(text, config, regexes):
    whitelist_rx, name_rx, last_rx = regexes
    placeholders = {}
    def stash(m):
        ph = f"<<WL{len(placeholders)}>>"
        placeholders[ph] = m.group(0)
        return ph
    protected = whitelist_rx.sub(stash, text)
    n = 0
    def replace_first(m):
        nonlocal n
        n += 1
        return config.replacement_token
    redacted = name_rx.sub(replace_first, protected)
    def replace_last(m):
        nonlocal n
        n += 1
        return config.last_name_replacement_token
    redacted = last_rx.sub(replace_last, redacted)
    for ph, orig in placeholders.items():
        redacted = redacted.replace(ph, orig)
    return redacted, n


def scrub_to_instructor_only(
    raw_path, out_path, config, *, report_path: Path | None = None
) -> RedactionStats:
    regexes = _build_regexes(config)
    student_dropped = 0
    instructor_kept = 0
    headers_kept = 0
    outcomes_kept = 0
    instr_redactions = 0
    outcome_redactions = 0

    out_path.parent.mkdir(parents=True, exist_ok=True)
    tmp = out_path.with_suffix(out_path.suffix + ".tmp")
    with open(raw_path, encoding="utf-8") as src, open(tmp, "w", encoding="utf-8") as dst:
        for line in src:
            line = line.strip()
            if not line:
                continue
            try:
                rec = json.loads(line)
            except json.JSONDecodeError:
                continue
            rt = rec.get("record_type", "")
            if rt == "session_header":
                headers_kept += 1
                dst.write(json.dumps(rec, ensure_ascii=False) + "\n")
            elif rt == "turn":
                speaker = rec.get("speaker", "")
                if speaker == "Anonymous Student":
                    student_dropped += 1
                    continue
                text, n = _redact_text(rec.get("text", "") or "", config, regexes)
                rec["text"] = text
                if n > 0:
                    rec["_redactions"] = n
                    instr_redactions += n
                instructor_kept += 1
                dst.write(json.dumps(rec, ensure_ascii=False) + "\n")
            elif rt == "problem_outcome":
                blob = json.dumps(rec, ensure_ascii=False)
                redacted_blob, n = _redact_text(blob, config, regexes)
                if n > 0:
                    rec = json.loads(redacted_blob)
                    rec["_redactions"] = n
                    outcome_redactions += n
                outcomes_kept += 1
                dst.write(json.dumps(rec, ensure_ascii=False) + "\n")
    os.replace(tmp, out_path)

    if report_path is None:
        report_path = out_path.with_name(out_path.stem + ".redaction-report.txt")
    report_body = (
        f"REDACTION REPORT (policy v{config.policy_version})\n"
        + "=" * 50 + "\n"
        + f"Source: {raw_path}\nOutput: {out_path}\n\n"
        + f"Dropped student turns: {student_dropped}\n"
        + f"Kept instructor turns: {instructor_kept}\n"
        + f"Kept session headers: {headers_kept}\n"
        + f"Kept problem outcomes: {outcomes_kept}\n"
        + f"Redactions in instructor turns: {instr_redactions}\n"
        + f"Redactions in problem outcomes: {outcome_redactions}\n"
    )
    report_path.parent.mkdir(parents=True, exist_ok=True)
    report_tmp = report_path.with_suffix(report_path.suffix + ".tmp")
    report_tmp.write_text(report_body, encoding="utf-8")
    os.replace(report_tmp, report_path)
    return RedactionStats(
        student_turns_dropped=student_dropped,
        instructor_turns_kept=instructor_kept,
        session_headers_kept=headers_kept,
        problem_outcomes_kept=outcomes_kept,
        name_redactions_in_instructor=instr_redactions,
        name_redactions_in_problem_outcomes=outcome_redactions,
    )