src/reason_interventions.py

"""
Sentence edits for reason faithfulness tests (UI + 07_faithfulness.py).

Each reason type has one targeted edit; re-predict after edit to see if hardness drops.
"""

from __future__ import annotations

import os
from dataclasses import dataclass

from syntax_complexity import find_hardest_span
from utils import REASON_ORDER


@dataclass
class InterventionResult:
    reason: str
    original_sentence: str
    edited_sentence: str
    edit_method: str  # "openai", "gemini", or "rule"


def _call_llm(prompt: str) -> tuple[str | None, str]:
    if os.environ.get("OPENAI_API_KEY"):
        try:
            from openai import OpenAI

            client = OpenAI()
            resp = client.chat.completions.create(
                model=os.environ.get("FAITHFULNESS_LLM", "gpt-4o-mini"),
                messages=[{"role": "user", "content": prompt}],
                temperature=0,
                max_tokens=256,
            )
            return resp.choices[0].message.content.strip(), "openai"
        except Exception:
            pass
    if os.environ.get("GOOGLE_API_KEY"):
        try:
            import google.generativeai as genai

            genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
            model = genai.GenerativeModel(os.environ.get("FAITHFULNESS_GEMINI", "gemini-1.5-flash"))
            return model.generate_content(prompt).text.strip(), "gemini"
        except Exception:
            pass
    return None, "rule"


def _rule_synonym_swap(sentence: str, target_word: str) -> str:
    simple = {
        "frankincense": "incense",
        "astonishment": "surprise",
        "partiality": "bias",
        "dominion": "rule",
        "assemblies": "meetings",
        "ewe": "sheep",
        "scribe": "writer",
        "inflammation": "swelling",
        "treaty": "deal",
    }
    rep = simple.get(target_word.lower(), target_word)
    if target_word in sentence:
        return sentence.replace(target_word, rep, 1)
    return sentence


def edit_lexical_rarity(sentence: str, target_word: str, use_llm: bool = True) -> InterventionResult:
    prompt = (
        f"Replace ONLY the word '{target_word}' in this sentence with a common, simple synonym "
        f"of the same meaning and part of speech. Change nothing else.\n"
        f"Sentence: {sentence}\n"
        f"Output only the edited sentence."
    )
    method = "rule"
    edited = _rule_synonym_swap(sentence, target_word)
    if use_llm:
        llm_out, method = _call_llm(prompt)
        if llm_out:
            edited = llm_out
    return InterventionResult("Lexical Rarity", sentence, edited, method)


def edit_contextual_ambiguity(sentence: str, target_word: str, use_llm: bool = True) -> InterventionResult:
    prompt = (
        f"Add at most 3 words to this sentence to make the meaning of '{target_word}' clear. "
        f"Do NOT change or remove '{target_word}'.\n"
        f"Sentence: {sentence}\n"
        f"Output only the edited sentence."
    )
    method = "rule"
    edited = f"{sentence} (meaning: {target_word})"
    if use_llm:
        llm_out, method = _call_llm(prompt)
        if llm_out:
            edited = llm_out
    return InterventionResult("Contextual Ambiguity", sentence, edited, method)


def edit_syntactic_complexity(sentence: str, target_word: str, use_llm: bool = True) -> InterventionResult:
    span, start, end = find_hardest_span(sentence)
    method = "rule"
    simplified = span
    if use_llm and span.strip() != sentence.strip():
        prompt = (
            f"Simplify ONLY this phrase to plain English. Keep all key words including '{target_word}'.\n"
            f"Phrase: {span}\n"
            f"Output only the simplified phrase."
        )
        llm_out, method = _call_llm(prompt)
        if llm_out:
            simplified = llm_out
    edited = sentence[:start] + simplified + sentence[end:]
    return InterventionResult("Syntactic Complexity", sentence, edited, method)


EDIT_FNS = {
    "Lexical Rarity": edit_lexical_rarity,
    "Contextual Ambiguity": edit_contextual_ambiguity,
    "Syntactic Complexity": edit_syntactic_complexity,
}


def apply_intervention(reason: str, sentence: str, target_word: str, use_llm: bool = True) -> InterventionResult:
    if reason not in EDIT_FNS:
        raise ValueError(f"Unknown reason: {reason}. Choose from {REASON_ORDER}")
    return EDIT_FNS[reason](sentence, target_word, use_llm=use_llm)


def apply_all_interventions(sentence: str, target_word: str, use_llm: bool = True) -> list[InterventionResult]:
    return [fn(sentence, target_word, use_llm=use_llm) for fn in EDIT_FNS.values()]