package ch.phoenixtechnologies.lc4j.workshop.guardrails;

import dev.langchain4j.data.message.AiMessage;
import io.quarkiverse.langchain4j.guardrails.OutputGuardrail;
import io.quarkiverse.langchain4j.guardrails.OutputGuardrailResult;
import io.quarkus.logging.Log;
import jakarta.enterprise.context.ApplicationScoped;

@ApplicationScoped
public class NumericOutputSanitizerGuard implements OutputGuardrail {
    @Override
    public OutputGuardrailResult validate(AiMessage responseFromLLM) {
        String llmResponse = responseFromLLM.text();

        try {
            double number = Double.parseDouble(llmResponse);
            return successWith(llmResponse, number);
        } catch (NumberFormatException e) {
            // ignore
        }

        Log.debugf("LLM output for expected numeric result: %s", llmResponse);

        String extractedNumber = extractNumber(llmResponse);
        if (extractedNumber != null) {
            Log.infof("Extracted number: %s", extractedNumber);
            try {
                double number = Double.parseDouble(extractedNumber);
                return successWith(extractedNumber, number);
            } catch (NumberFormatException e) {
                // ignore
            }
        }

        return failure("Unable to extract a number from LLM response: " + llmResponse);
    }

    private String extractNumber(String text) {
        int lastDigitPosition = text.length() - 1;
        while (lastDigitPosition >= 0) {
            if (Character.isDigit(text.charAt(lastDigitPosition))) {
                break;
            }
            lastDigitPosition--;
        }
        if (lastDigitPosition < 0) {
            return null;
        }
        int numberBegin = lastDigitPosition;
        while (numberBegin >= 0) {
            if (!Character.isDigit(text.charAt(numberBegin)) && text.charAt(numberBegin) != '.') {
                break;
            }
            numberBegin--;
        }
        return text.substring(numberBegin + 1, lastDigitPosition + 1);
    }
}
