package dev.sonpx.loyalty.agent.service;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import java.util.Optional;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

@Slf4j
@Component
@RequiredArgsConstructor
public class ToolResultPresentationProcessor {

    private static final String INTERNAL_INSTRUCTION_FIELD = "_agent_instruction";
    private static final String AGENT_INSTRUCTION_FIELD = "agent_instruction";

    private final ObjectMapper objectMapper;

    public String process(String result) {
        if (!StringUtils.hasText(result)) {
            return result;
        }

        try {
            JsonNode root = objectMapper.readTree(result);
            return root.isArray() ? processMcpContentArray(root, result) : processJsonObject(root, result);
        } catch (JsonProcessingException e) {
            log.debug("Tool result is not JSON.");
            return result;
        }
    }

    private String processMcpContentArray(JsonNode root, String originalResult)
            throws JsonProcessingException {
        boolean modified = false;

        for (JsonNode contentItem : root) {
            if (!(contentItem instanceof ObjectNode objectContentItem)) {
                continue;
            }

            JsonNode textNode = objectContentItem.get("text");
            if (textNode == null || !textNode.isTextual()) {
                continue;
            }

            Optional<JsonNode> parsedText = readJson(textNode.asText());
            if (parsedText.isEmpty()) {
                continue;
            }

            JsonNode parsedNode = parsedText.get();
            if (processInstruction(parsedNode)) {
                modified = true;
                objectContentItem.put("text", objectMapper.writeValueAsString(parsedNode));
            }
        }

        return modified ? objectMapper.writeValueAsString(root) : originalResult;
    }

    private String processJsonObject(JsonNode root, String originalResult) throws JsonProcessingException {
        if (processInstruction(root)) {
            return objectMapper.writeValueAsString(root);
        }
        return originalResult;
    }

    private Optional<JsonNode> readJson(String content) {
        try {
            return Optional.of(objectMapper.readTree(content));
        } catch (JsonProcessingException e) {
            return Optional.empty();
        }
    }

    private boolean processInstruction(JsonNode node) {
        if (!(node instanceof ObjectNode objectNode)) {
            return false;
        }

        JsonNode instructionNode = objectNode.get(INTERNAL_INSTRUCTION_FIELD);
        if (instructionNode == null || !instructionNode.isTextual()) {
            return false;
        }

        String instruction = instructionNode.asText();
        if (!StringUtils.hasText(instruction)) {
            return false;
        }

        objectNode.remove(INTERNAL_INSTRUCTION_FIELD);
        objectNode.put(AGENT_INSTRUCTION_FIELD, instruction);
        return true;
    }
}
