"""
Triton Python backend for RAI Image Moderation v1 (ShieldGemma2).

This wraps the SageMaker inference.py logic into Triton Python backend format.
Input:  image (BYTES) - base64-encoded image (optionally with custom_policies in JSON)
Output: result (BYTES) - JSON list of policy predictions matching SageMaker output format:
        [{"policy": "...", "prediction": {"Unsafe": 0.9, "Safe": 0.1}}, ...]
"""

import json
import base64
import os
import numpy as np
import triton_python_backend_utils as pb_utils


# Default policies (matching the original SageMaker inference.py)
DEFAULT_POLICIES = {
    "Hate/Discrimination": "The image must not contain content that promotes attacks based on personal characteristics or content that claims certain protected groups such as gender, race, religion, sexual orientation, etc. are better than others.",
    "Misinformation": "The image must not contain content seeking to spread misinformation, disinformation, or conspiracy theories that undermine public trust.",
    "Jailbreak/Prompt Injection": "The image must not contain content attempting to manipulate AI models to disobey explicit instructions or safety policies.",
    "Personally Identifiable Information (PII)": "The image must not contain content that aims to extract sensitive personal information such as social security numbers, credit card information, driver licenses, passport numbers, and medical information.",
    "Politics": "The image must not contain content involving political campaigning or lobbying that may interfere with electoral or democratic processes.",
    "Impersonation": "The image must not contain content that impersonates specific individuals, fictional characters, or organizations.",
    "Specialist Advice": "The image must not provide advice that would ordinarily be provided by a qualified professional including legal, medical, financial advice.",
    "High-risk Decisions": "The image must not contain content involving high stakes decision making that may affect an individual's rights, safety, health or well-being.",
}


class TritonPythonModel:
    """Triton Python backend model class for ShieldGemma2 image safety classification."""

    def initialize(self, args):
        """
        Called once when the model is loaded.
        """
        import torch
        from transformers import AutoProcessor, ShieldGemma2ForImageClassification

        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        pb_utils.Logger.log_info(f"Loading ShieldGemma2 model on device: {self.device}")

        # Model artifacts are in the parent directory (alongside config.pbtxt)
        model_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

        self.model = ShieldGemma2ForImageClassification.from_pretrained(
            model_dir,
            torch_dtype=torch.bfloat16
        ).to(self.device).eval()

        self.processor = AutoProcessor.from_pretrained(model_dir)
        pb_utils.Logger.log_info("ShieldGemma2 model loaded successfully")

    def execute(self, requests):
        """
        Called for each inference request batch.
        """
        import torch
        from PIL import Image
        from io import BytesIO

        responses = []

        for request in requests:
            # Extract base64-encoded image from input tensor
            image_tensor = pb_utils.get_input_tensor_by_name(request, "image")
            input_bytes = image_tensor.as_numpy()[0].decode("utf-8")

            # Input can be either a raw base64 string or a JSON payload
            try:
                input_data = json.loads(input_bytes)
                if isinstance(input_data, dict):
                    image_b64 = input_data.get('image', input_bytes)
                    custom_policies = input_data.get('custom_policies', None)
                else:
                    image_b64 = input_bytes
                    custom_policies = None
            except (json.JSONDecodeError, TypeError):
                image_b64 = input_bytes
                custom_policies = None

            # Decode image
            img_bytes = base64.b64decode(image_b64)
            pil_image = Image.open(BytesIO(img_bytes)).convert('RGB')

            # Use client-provided policies or defaults
            policies_dict = custom_policies if custom_policies else DEFAULT_POLICIES

            # Build key mapping (lowercase + underscore)
            custom_policies_dict = {}
            policy_name_mapping = {}
            for policy_name, policy_text in policies_dict.items():
                key = policy_name.lower()
                for ch in ['/', ' ', '(', ')', '-']:
                    key = key.replace(ch, '_')
                key = key.replace('__', '_').strip('_')
                custom_policies_dict[key] = policy_text
                policy_name_mapping[key] = policy_name

            policies_to_evaluate = list(custom_policies_dict.keys())

            # Run inference
            model_inputs = self.processor(
                images=[pil_image],
                custom_policies=custom_policies_dict,
                policies=policies_to_evaluate,
                return_tensors="pt"
            ).to(self.device)

            with torch.inference_mode():
                output = self.model(**model_inputs)

            probabilities = output.probabilities  # Shape: [num_policies, 2]

            # Build results list matching SageMaker output format
            results = []
            for idx, policy_key in enumerate(policies_to_evaluate):
                unsafe_prob = 1.0 - float(probabilities[idx][1])
                safe_prob = float(probabilities[idx][1])
                results.append({
                    "policy": policy_name_mapping[policy_key],
                    "prediction": {"Unsafe": unsafe_prob, "Safe": safe_prob}
                })

            # Encode result as JSON bytes
            result_json = json.dumps(results).encode('utf-8')
            out_result = pb_utils.Tensor("result", np.array([result_json], dtype=object))

            responses.append(pb_utils.InferenceResponse(output_tensors=[out_result]))

        return responses

    def finalize(self):
        """Called when model is unloaded."""
        pb_utils.Logger.log_info("ShieldGemma2 model unloaded")