#!/usr/bin/env python3
"""
Handle Claude Fable 5 refusals on Amazon Bedrock (Converse API).

Fable 5 returns HTTP 200 with stop_reason="refusal" when blocking classifiers
fire (cybersecurity, biology, etc.). Treat this as a primary response path,
not an exception. Prompt-stage refusals are not billed; mid-stream refusals
bill tokens generated before the block.

Requires: boto3, AWS credentials with bedrock:InvokeModel on anthropic.claude-fable-5
Region: us-east-1 (confirm in model card before production)
"""
from __future__ import annotations

import json
import sys
from typing import Any

import boto3
from botocore.exceptions import ClientError

MODEL_ID = "global.anthropic.claude-fable-5"  # or us.anthropic.claude-fable-5 for geo
REGION = "us-east-1"
FALLBACK_MODEL_ID = "us.anthropic.claude-sonnet-4-6"  # adjust to your approved fallback


def converse(client: Any, model_id: str, user_text: str) -> dict[str, Any]:
    return client.converse(
        modelId=model_id,
        messages=[{"role": "user", "content": [{"text": user_text}]}],
        inferenceConfig={"maxTokens": 4096},
    )


def extract_stop_reason(response: dict[str, Any]) -> str | None:
    output = response.get("output", {})
    message = output.get("message", {})
    return message.get("stopReason") or message.get("stop_reason")


def extract_text(response: dict[str, Any]) -> str:
    content = response.get("output", {}).get("message", {}).get("content", [])
    for block in content:
        if "text" in block:
            return block["text"]
    return ""


def invoke_with_refusal_handling(client: Any, prompt: str) -> dict[str, Any]:
    response = converse(client, MODEL_ID, prompt)
    stop_reason = extract_stop_reason(response)

    if stop_reason == "refusal":
        # Log for compliance review; do not retry the same prompt blindly.
        print(
            json.dumps(
                {
                    "event": "fable_refusal",
                    "model": MODEL_ID,
                    "action": "fallback_or_human_review",
                }
            ),
            file=sys.stderr,
        )
        # Option A: escalate to human. Option B: bounded fallback model (policy-dependent).
        fallback = converse(client, FALLBACK_MODEL_ID, prompt)
        return {
            "primary": response,
            "fallback": fallback,
            "text": extract_text(fallback),
            "handled_refusal": True,
        }

    return {
        "primary": response,
        "text": extract_text(response),
        "handled_refusal": False,
    }


def main() -> None:
    client = boto3.client("bedrock-runtime", region_name=REGION)
    prompt = (
        "Summarize the indemnification clause risks in the attached MSA template "
        "for a B2B SaaS vendor."
    )
    try:
        result = invoke_with_refusal_handling(client, prompt)
        print(result["text"])
    except ClientError as exc:
        if exc.response["Error"]["Code"] == "ValidationException":
            # Common at launch: data retention not opted in
            print(
                "Enable provider_data_share via Bedrock Data Retention API before invoking Fable 5.",
                file=sys.stderr,
            )
        raise


if __name__ == "__main__":
    main()