"""Hard-negative mining for recursive OBLITERATUS runs.

This module turns refusal audits into reusable residue sets and weighted prompt
pairs.  The design intentionally separates *mining metadata* from *training
prompt reconstruction*: by default audit files can store only dataset keys,
indices, hashes, and refusal reasons; later runs reconstruct prompt text from
local registered corpora.
"""

from __future__ import annotations

import hashlib
import json
import re
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Iterable

from obliteratus.prompts import _generate_harmless_counterparts, load_dataset_source

_REFUSAL_REASON_TAGS: tuple[tuple[str, str], ...] = (
    ("terrorism", "terrorism|terrorist|dirty bomb|biological terrorism"),
    ("illegal_crime", "illegal|against the law|crime|criminal|felony|violation"),
    ("privacy_intrusion", "privacy|surveillance|monitor|smart home|spy"),
    ("sabotage", "sabotage|industrial equipment|machinery|critical infrastructure"),
    ("safety_banner", "dangerous|red flag|stop|do not|don't|never attempt"),
    ("generic_refusal", "i can't|i cannot|can't assist|cannot assist|unable to"),
)


@dataclass(frozen=True)
class ResidueExample:
    """A single mined hard-negative prompt reference."""

    dataset: str
    index: int  # zero-based index into the dataset's harmful prompt list
    reason: str = ""
    prompt_sha256_12: str = ""
    tier: str | None = None
    weight: int = 1
    harmless: str | None = None
    tags: tuple[str, ...] = ()

    def to_json(self) -> dict[str, Any]:
        data = asdict(self)
        data["tags"] = list(self.tags)
        return data


def prompt_hash(text: str, n: int = 12) -> str:
    """Return a short SHA-256 fingerprint for a prompt."""

    return hashlib.sha256(text.encode("utf-8")).hexdigest()[:n]


def classify_refusal_reason(reason: str, preview: str = "") -> tuple[str, ...]:
    """Map a refusal reason/preview into stable residue tags."""

    haystack = f"{reason}\n{preview}".lower()
    tags: list[str] = []
    for tag, pattern in _REFUSAL_REASON_TAGS:
        if re.search(pattern, haystack):
            tags.append(tag)
    return tuple(tags or ["other_refusal"])


def load_residue_file(path: str | Path) -> list[ResidueExample]:
    """Load a residue JSON file.

    Supported shapes:
      - {"examples": [...]} produced by ``save_residue_file``
      - {"refusals": [...]} audit output from scripts/qwen36_refusal_audit.py
      - {"model": {"refusals": [...]}} combined audit summary
      - raw list of residue/example dicts
    """

    raw = json.loads(Path(path).read_text())
    rows: list[dict[str, Any]] = []
    if isinstance(raw, list):
        rows = [x for x in raw if isinstance(x, dict)]
    elif isinstance(raw, dict) and "examples" in raw:
        rows = [x for x in raw["examples"] if isinstance(x, dict)]
    elif isinstance(raw, dict) and "refusals" in raw:
        rows = [x for x in raw["refusals"] if isinstance(x, dict)]
    elif isinstance(raw, dict):
        for value in raw.values():
            if isinstance(value, dict) and isinstance(value.get("refusals"), list):
                rows.extend(x for x in value["refusals"] if isinstance(x, dict))

    examples: list[ResidueExample] = []
    for row in rows:
        dataset = str(row.get("dataset") or row.get("dataset_key") or "builtin")
        if "index" in row:
            index = int(row["index"])
        elif "global_index_1based" in row:
            index = int(row["global_index_1based"]) - 1
        else:
            continue
        reason = str(row.get("reason", ""))
        preview = str(row.get("response_preview", ""))
        tags_raw = row.get("tags")
        tags = tuple(tags_raw) if isinstance(tags_raw, list) else classify_refusal_reason(reason, preview)
        examples.append(
            ResidueExample(
                dataset=dataset,
                index=index,
                reason=reason,
                prompt_sha256_12=str(row.get("prompt_sha256_12") or ""),
                tier=row.get("tier"),
                weight=max(1, int(row.get("weight", 1))),
                harmless=row.get("harmless"),
                tags=tags,
            )
        )
    return dedupe_residue(examples)


def save_residue_file(examples: Iterable[ResidueExample], path: str | Path) -> Path:
    """Write canonical residue JSON without prompt text."""

    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    rows = [ex.to_json() for ex in dedupe_residue(examples)]
    path.write_text(json.dumps({"version": 1, "examples": rows}, indent=2))
    return path


def dedupe_residue(examples: Iterable[ResidueExample]) -> list[ResidueExample]:
    """Dedupe by dataset/index while preserving max weight and merged tags."""

    merged: dict[tuple[str, int], ResidueExample] = {}
    for ex in examples:
        key = (ex.dataset, ex.index)
        prev = merged.get(key)
        if prev is None:
            merged[key] = ex
            continue
        tags = tuple(sorted(set(prev.tags) | set(ex.tags)))
        merged[key] = ResidueExample(
            dataset=ex.dataset,
            index=ex.index,
            reason=ex.reason or prev.reason,
            prompt_sha256_12=ex.prompt_sha256_12 or prev.prompt_sha256_12,
            tier=ex.tier or prev.tier,
            weight=max(prev.weight, ex.weight),
            harmless=ex.harmless or prev.harmless,
            tags=tags,
        )
    return list(merged.values())


def build_weighted_prompt_pairs(
    base_dataset: str = "builtin",
    residue_files: Iterable[str | Path] = (),
    residue_weight: int = 5,
    max_residue: int | None = None,
) -> tuple[list[str], list[str], dict[str, Any]]:
    """Load base prompt pairs and append repeated hard-negative residue pairs.

    The returned metadata is safe for logs/model cards: it records counts,
    indices, hashes, and tags, but not prompt text.
    """

    harmful, harmless = load_dataset_source(base_dataset)
    base_n = min(len(harmful), len(harmless))
    harmful = list(harmful[:base_n])
    harmless = list(harmless[:base_n])

    examples: list[ResidueExample] = []
    for path in residue_files:
        examples.extend(load_residue_file(path))
    examples = dedupe_residue(examples)
    if max_residue is not None:
        examples = examples[:max_residue]

    added = 0
    records: list[dict[str, Any]] = []
    for ex in examples:
        ds_harmful, ds_harmless = load_dataset_source(ex.dataset)
        if ex.index < 0 or ex.index >= len(ds_harmful):
            continue
        prompt = ds_harmful[ex.index]
        if ex.prompt_sha256_12 and prompt_hash(prompt) != ex.prompt_sha256_12:
            # The corpus changed under the audit. Skip rather than poisoning the run.
            continue
        counterpart = ex.harmless
        if not counterpart:
            if ex.index < len(ds_harmless):
                counterpart = ds_harmless[ex.index]
            else:
                counterpart = _generate_harmless_counterparts(1)[0]
        copies = max(1, residue_weight * ex.weight)
        harmful.extend([prompt] * copies)
        harmless.extend([counterpart] * copies)
        added += copies
        rec = ex.to_json()
        rec["effective_copies"] = copies
        rec["prompt_sha256_12"] = rec["prompt_sha256_12"] or prompt_hash(prompt)
        records.append(rec)

    meta = {
        "base_dataset": base_dataset,
        "base_pairs": base_n,
        "residue_examples": len(records),
        "residue_added_pairs": added,
        "total_pairs": len(harmful),
        "residue_records": records,
    }
    return harmful, harmless, meta