package com.ankurm.imagegen;

import java.util.ArrayList;
import java.util.List;

import org.springframework.ai.moderation.Categories;
import org.springframework.ai.moderation.ModerationModel;
import org.springframework.ai.moderation.ModerationPrompt;
import org.springframework.ai.moderation.ModerationResult;

/**
 * The moderation step Spring AI does not run for you. It asks a {@link ModerationModel} about the
 * prompt before any image is requested and fails closed: if the moderation call itself throws, the
 * prompt is refused rather than waved through.
 *
 * <p>This inspects the TEXT. It says nothing about the picture that comes back.
 */
public final class PromptGate {

    private final ModerationModel moderation;
    private final ImagePolicy policy;

    public PromptGate(ModerationModel moderation, ImagePolicy policy) {
        this.moderation = moderation;
        this.policy = policy;
    }

    public void check(String prompt) {
        if (prompt == null || prompt.isBlank()) {
            throw new RejectedException("empty-prompt", "prompt is blank");
        }
        if (prompt.length() > policy.maxPromptChars()) {
            throw new RejectedException("prompt-too-long", prompt.length() + " > " + policy.maxPromptChars());
        }
        if (prompt.chars().anyMatch(c -> Character.isISOControl(c) && c != '\n' && c != '\t')) {
            throw new RejectedException("control-characters", "prompt contains control characters");
        }
        List<String> flagged = new ArrayList<>();
        try {
            for (ModerationResult r : moderation.call(new ModerationPrompt(prompt)).getResult().getOutput().getResults()) {
                if (r.isFlagged()) {
                    flagged.addAll(flaggedCategories(r.getCategories()));
                }
            }
        } catch (RuntimeException e) {
            throw new RejectedException("moderation-unavailable", e.getClass().getSimpleName());
        }
        if (!flagged.isEmpty()) {
            throw new RejectedException("prompt-flagged", String.join(",", flagged));
        }
    }

    private static List<String> flaggedCategories(Categories c) {
        List<String> out = new ArrayList<>();
        if (c == null) return List.of("unspecified");
        if (c.isSexual()) out.add("sexual");
        if (c.isSexualMinors()) out.add("sexual/minors");
        if (c.isHate()) out.add("hate");
        if (c.isHarassment()) out.add("harassment");
        if (c.isSelfHarm()) out.add("self-harm");
        if (c.isViolence()) out.add("violence");
        if (c.isViolenceGraphic()) out.add("violence/graphic");
        return out.isEmpty() ? List.of("unspecified") : out;
    }
}