/*
 * Ollama4j - Java library for interacting with Ollama server.
 * Copyright (c) 2025 Amith Koujalgi and contributors.
 *
 * Licensed under the MIT License (the "License");
 * you may not use this file except in compliance with the License.
 *
*/
package io.github.ollama4j.tools;

import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.core.JsonProcessingException;
import io.github.ollama4j.utils.Utils;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;

public class Tools {
    @Data
    @Builder
    public static class ToolSpecification {
        private String functionName;
        private String functionDescription;
        private PromptFuncDefinition toolPrompt;
        private ToolFunction toolFunction;
    }

    @Data
    @JsonIgnoreProperties(ignoreUnknown = true)
    @Builder
    @NoArgsConstructor
    @AllArgsConstructor
    public static class PromptFuncDefinition {
        private String type;
        private PromptFuncSpec function;

        @Data
        @Builder
        @NoArgsConstructor
        @AllArgsConstructor
        public static class PromptFuncSpec {
            private String name;
            private String description;
            private Parameters parameters;
        }

        @Data
        @Builder
        @NoArgsConstructor
        @AllArgsConstructor
        public static class Parameters {
            private String type;
            private Map<String, Property> properties;
            private List<String> required;
        }

        @Data
        @Builder
        @NoArgsConstructor
        @AllArgsConstructor
        public static class Property {
            private String type;
            private String description;

            @JsonProperty("enum")
            @JsonInclude(JsonInclude.Include.NON_NULL)
            private List<String> enumValues;

            @JsonIgnore private boolean required;
        }
    }

    public static class PropsBuilder {
        private final Map<String, PromptFuncDefinition.Property> props = new HashMap<>();

        public PropsBuilder withProperty(String key, PromptFuncDefinition.Property property) {
            props.put(key, property);
            return this;
        }

        public Map<String, PromptFuncDefinition.Property> build() {
            return props;
        }
    }

    public static class PromptBuilder {
        private final List<PromptFuncDefinition> tools = new ArrayList<>();

        private String promptText;

        public String build() throws JsonProcessingException {
            return "[AVAILABLE_TOOLS] "
                    + Utils.getObjectMapper().writeValueAsString(tools)
                    + "[/AVAILABLE_TOOLS][INST] "
                    + promptText
                    + " [/INST]";
        }

        public PromptBuilder withPrompt(String prompt) throws JsonProcessingException {
            promptText = prompt;
            return this;
        }

        public PromptBuilder withToolSpecification(ToolSpecification spec) {
            PromptFuncDefinition def = new PromptFuncDefinition();
            def.setType("function");

            PromptFuncDefinition.PromptFuncSpec functionDetail =
                    new PromptFuncDefinition.PromptFuncSpec();
            functionDetail.setName(spec.getFunctionName());
            functionDetail.setDescription(spec.getFunctionDescription());

            PromptFuncDefinition.Parameters parameters = new PromptFuncDefinition.Parameters();
            parameters.setType("object");
            parameters.setProperties(spec.getToolPrompt().getFunction().parameters.getProperties());

            List<String> requiredValues = new ArrayList<>();
            for (Map.Entry<String, PromptFuncDefinition.Property> p :
                    spec.getToolPrompt().getFunction().getParameters().getProperties().entrySet()) {
                if (p.getValue().isRequired()) {
                    requiredValues.add(p.getKey());
                }
            }
            parameters.setRequired(requiredValues);
            functionDetail.setParameters(parameters);
            def.setFunction(functionDetail);

            tools.add(def);
            return this;
        }
    }
}
