/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.flink.agents.integration.test;

import org.apache.flink.agents.api.Event;
import org.apache.flink.agents.api.InputEvent;
import org.apache.flink.agents.api.OutputEvent;
import org.apache.flink.agents.api.agents.Agent;
import org.apache.flink.agents.api.annotation.Action;
import org.apache.flink.agents.api.annotation.ChatModelConnection;
import org.apache.flink.agents.api.annotation.ChatModelSetup;
import org.apache.flink.agents.api.annotation.Tool;
import org.apache.flink.agents.api.annotation.ToolParam;
import org.apache.flink.agents.api.chat.messages.ChatMessage;
import org.apache.flink.agents.api.chat.messages.MessageRole;
import org.apache.flink.agents.api.chat.model.BaseChatModelConnection;
import org.apache.flink.agents.api.chat.model.BaseChatModelSetup;
import org.apache.flink.agents.api.context.RunnerContext;
import org.apache.flink.agents.api.event.ChatRequestEvent;
import org.apache.flink.agents.api.event.ChatResponseEvent;
import org.apache.flink.agents.api.resource.ResourceDescriptor;
import org.apache.flink.agents.api.resource.ResourceName;
import org.apache.flink.agents.api.resource.ResourceType;

import java.util.Collections;
import java.util.List;

/**
 * Agent example that integrates an external Ollama chat model into Flink Agents.
 *
 * <p>This class demonstrates how to:
 *
 * <ul>
 *   <li>Declare a chat model connection using {@link ChatModelConnection} metadata pointing to an
 *       implementation of {@link BaseChatModelConnection}
 *   <li>Declare a chat model setup using {@link ChatModelSetup} metadata pointing to an
 *       implementation of {@link BaseChatModelSetup}
 *   <li>Expose callable tools via {@link Tool} annotated static methods (temperature conversion,
 *       BMI, random number)
 *   <li>Fetch a chat model from the {@link RunnerContext} and perform a single-turn chat
 *   <li>Emit the model response as an {@link OutputEvent}
 * </ul>
 *
 * <p>The {@code ollamaChatModel()} method publishes a resource with type {@link
 * ResourceType#CHAT_MODEL} so it can be retrieved at runtime inside the {@code process} action. The
 * resource is configured with the connection name, the model name and the list of tool names that
 * the model is allowed to call.
 */
public class ChatModelIntegrationAgent extends Agent {
    public static final String OLLAMA_MODEL = "qwen3:0.6b";

    @ChatModelConnection
    public static ResourceDescriptor chatModelConnection() {
        String provider = System.getProperty("MODEL_PROVIDER", "OLLAMA");
        if (provider.equals("OLLAMA")) {
            return ResourceDescriptor.Builder.newBuilder(ResourceName.ChatModel.OLLAMA_CONNECTION)
                    .addInitialArgument("endpoint", "http://localhost:11434")
                    .addInitialArgument("requestTimeout", 240)
                    .build();
        } else if (provider.equals("AZURE")) {
            String endpoint = System.getenv().get("AZURE_ENDPOINT");
            String apiKey = System.getenv().get("AZURE_API_KEY");
            return ResourceDescriptor.Builder.newBuilder(ResourceName.ChatModel.AZURE_CONNECTION)
                    .addInitialArgument("endpoint", endpoint)
                    .addInitialArgument("apiKey", apiKey)
                    .build();
        } else if (provider.equals("OPENAI")) {
            String apiKey = System.getenv().get("OPENAI_API_KEY");
            return ResourceDescriptor.Builder.newBuilder(
                            ResourceName.ChatModel.OPENAI_COMPLETIONS_CONNECTION)
                    .addInitialArgument("api_key", apiKey)
                    .build();
        } else if (provider.equals("OPENAI_RESPONSES")) {
            return ResourceDescriptor.Builder.newBuilder(
                            ResourceName.ChatModel.OPENAI_RESPONSES_CONNECTION)
                    .addInitialArgument("api_key", System.getenv().get("OPENAI_API_KEY"))
                    .build();
        } else if (provider.equals("AZURE_OPENAI")) {
            return ResourceDescriptor.Builder.newBuilder(
                            ResourceName.ChatModel.AZURE_OPENAI_CONNECTION)
                    .addInitialArgument("api_key", System.getenv().get("AZURE_OPENAI_API_KEY"))
                    .addInitialArgument(
                            "api_version", System.getenv().get("AZURE_OPENAI_API_VERSION"))
                    .addInitialArgument(
                            "azure_endpoint", System.getenv().get("AZURE_OPENAI_ENDPOINT"))
                    .addInitialArgument("azure_url_path_mode", "LEGACY")
                    .build();
        } else if (provider.equals("ANTHROPIC")) {
            String apiKey = System.getenv().get("ANTHROPIC_API_KEY");
            return ResourceDescriptor.Builder.newBuilder(
                            ResourceName.ChatModel.ANTHROPIC_CONNECTION)
                    .addInitialArgument("api_key", apiKey)
                    .addInitialArgument("timeout", 240)
                    .build();
        } else {
            throw new RuntimeException(String.format("Unknown model provider %s", provider));
        }
    }

    @ChatModelSetup
    public static ResourceDescriptor chatModel() {
        String provider = System.getProperty("MODEL_PROVIDER", "OLLAMA");

        if (provider.equals("OLLAMA")) {
            return ResourceDescriptor.Builder.newBuilder(ResourceName.ChatModel.OLLAMA_SETUP)
                    .addInitialArgument("connection", "chatModelConnection")
                    .addInitialArgument("model", OLLAMA_MODEL)
                    .addInitialArgument("think", false)
                    .addInitialArgument(
                            "tools",
                            List.of("calculateBMI", "convertTemperature", "createRandomNumber"))
                    .build();
        } else if (provider.equals("AZURE")) {
            return ResourceDescriptor.Builder.newBuilder(ResourceName.ChatModel.AZURE_SETUP)
                    .addInitialArgument("connection", "chatModelConnection")
                    .addInitialArgument("model", "gpt-4o")
                    .addInitialArgument(
                            "tools",
                            List.of("calculateBMI", "convertTemperature", "createRandomNumber"))
                    .build();
        } else if (provider.equals("ANTHROPIC")) {
            return ResourceDescriptor.Builder.newBuilder(ResourceName.ChatModel.ANTHROPIC_SETUP)
                    .addInitialArgument("connection", "chatModelConnection")
                    .addInitialArgument("model", "claude-sonnet-4-20250514")
                    .addInitialArgument(
                            "tools",
                            List.of("calculateBMI", "convertTemperature", "createRandomNumber"))
                    .build();
        } else if (provider.equals("OPENAI")) {
            return ResourceDescriptor.Builder.newBuilder(
                            ResourceName.ChatModel.OPENAI_COMPLETIONS_SETUP)
                    .addInitialArgument("connection", "chatModelConnection")
                    .addInitialArgument("model", "gpt-4o-mini")
                    .addInitialArgument(
                            "tools",
                            List.of("calculateBMI", "convertTemperature", "createRandomNumber"))
                    .build();
        } else if (provider.equals("OPENAI_RESPONSES")) {
            return ResourceDescriptor.Builder.newBuilder(
                            ResourceName.ChatModel.OPENAI_RESPONSES_SETUP)
                    .addInitialArgument("connection", "chatModelConnection")
                    .addInitialArgument("model", "gpt-4o-mini")
                    .addInitialArgument(
                            "tools",
                            List.of("calculateBMI", "convertTemperature", "createRandomNumber"))
                    .build();
        } else if (provider.equals("AZURE_OPENAI")) {
            return ResourceDescriptor.Builder.newBuilder(ResourceName.ChatModel.AZURE_OPENAI_SETUP)
                    .addInitialArgument("connection", "chatModelConnection")
                    .addInitialArgument("model", System.getenv().get("AZURE_OPENAI_DEPLOYMENT"))
                    .addInitialArgument(
                            "tools",
                            List.of("calculateBMI", "convertTemperature", "createRandomNumber"))
                    .build();
        } else {
            throw new RuntimeException(String.format("Unknown model provider %s", provider));
        }
    }

    @Tool(description = "Converts temperature between Celsius and Fahrenheit")
    public static double convertTemperature(
            @ToolParam(name = "value", description = "Temperature value to convert") Double value,
            @ToolParam(
                            name = "fromUnit",
                            description = "Source unit ('C' for Celsius or 'F' for Fahrenheit)")
                    String fromUnit,
            @ToolParam(
                            name = "toUnit",
                            description = "Target unit ('C' for Celsius or 'F' for Fahrenheit)")
                    String toUnit) {

        fromUnit = fromUnit.toUpperCase();
        toUnit = toUnit.toUpperCase();

        if (fromUnit.equals(toUnit)) {
            return value;
        }

        if (fromUnit.contains("C") && toUnit.contains("F")) {
            return (value * 9 / 5) + 32;
        } else if (fromUnit.contains("F") && toUnit.contains("C")) {
            return (value - 32) * 5 / 9;
        } else {
            throw new IllegalArgumentException("Invalid temperature units. Use 'C' or 'F'");
        }
    }

    @Tool(description = "Calculates Body Mass Index (BMI)")
    public static double calculateBMI(
            @ToolParam(name = "weightKg", description = "Weight in kilograms") Double weightKg,
            @ToolParam(name = "heightM", description = "Height in meters") Double heightM) {

        if (weightKg <= 0 || heightM <= 0) {
            throw new IllegalArgumentException("Weight and height must be positive values");
        }
        return weightKg / (heightM * heightM);
    }

    @Tool(description = "Create a random number")
    public static double createRandomNumber() {
        return Math.random();
    }

    @Action(listenEventTypes = {InputEvent.EVENT_TYPE})
    public static void process(Event event, RunnerContext ctx) throws Exception {
        InputEvent inputEvent = InputEvent.fromEvent(event);
        ctx.sendEvent(
                new ChatRequestEvent(
                        "chatModel",
                        Collections.singletonList(
                                new ChatMessage(
                                        MessageRole.USER, (String) inputEvent.getInput()))));
    }

    @Action(listenEventTypes = {ChatResponseEvent.EVENT_TYPE})
    public static void processChatResponse(Event event, RunnerContext ctx) {
        ChatResponseEvent chatResponse = ChatResponseEvent.fromEvent(event);
        ctx.sendEvent(new OutputEvent(chatResponse.getResponse().getContent()));
    }
}
