"""rewrite.py, M5: an engineered prompt that uses FEW-SHOT + STRUCTURED OUTPUT.

This is the "good" prompt from the A/B test, leveled up:
  - a clear SYSTEM role and rules,
  - a FEW-SHOT example (one worked example to lock the format and tone),
  - STRUCTURED OUTPUT: we ask for JSON and parse it into a Python dict (M3's json!),
    so the result is usable in code, not just text on screen.

Run (venv active, key in .env, from this folder):
    python rewrite.py
"""

import os
import json
from dotenv import load_dotenv
import anthropic

load_dotenv()
client = anthropic.Anthropic()
MODEL = "claude-opus-4-8"        # swap to "claude-haiku-4-5" for cheaper experiments

SYSTEM = (
    "You rewrite blunt or rushed messages into warm, professional ones. "
    "Keep the original meaning and any facts; be concise; remove blame. "
    'Return ONLY a JSON object with exactly these keys: '
    '"subject" (a short subject line), "body" (the rewritten message), '
    '"tone_note" (one sentence on what you changed).'
)

# FEW-SHOT: one example teaches the model the exact shape and tone we want.
EXAMPLES = [
    {"role": "user", "content": "rewrite: need the report today. you're late again."},
    {"role": "assistant", "content": json.dumps({
        "subject": "Quick check-in on the report",
        "body": "Hi! Following up on the report, could you send it over today? "
                "If anything's blocking you, let me know and I'll help.",
        "tone_note": "Turned the blame into a supportive, time-bound request.",
    })},
]


def rewrite(blunt_message):
    """Rewrite a blunt message; return a parsed dict with subject/body/tone_note."""
    messages = EXAMPLES + [{"role": "user", "content": f"rewrite: {blunt_message}"}]
    response = client.messages.create(
        model=MODEL, max_tokens=500, system=SYSTEM, messages=messages,
    )
    text = response.content[0].text.strip()
    # The model *usually* returns clean JSON, but may wrap it in ```json fences.
    # Strip them defensively, then parse. (M6 shows the API feature that GUARANTEES JSON.)
    if text.startswith("```"):
        text = text.removeprefix("```json").removeprefix("```").removesuffix("```").strip()
    return json.loads(text)


if __name__ == "__main__":
    blunt = input("Paste a blunt message to rewrite: ")
    try:
        result = rewrite(blunt)
        print(f"\nSubject: {result['subject']}\n")
        print(result["body"])
        print(f"\n(What changed: {result['tone_note']})")
    except json.JSONDecodeError:
        print("The model didn't return clean JSON that time, run it again, or see M6 for the "
              "API feature that guarantees JSON.")