/*
 * Copyright 2024-2026 the original author or authors.
 *
 * Licensed 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
 *
 *      https://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 com.alibaba.cloud.ai.graph.agent.scheduling;

import com.alibaba.cloud.ai.dashscope.api.DashScopeApi;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatModel;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.agent.ReactAgent;
import com.alibaba.cloud.ai.graph.scheduling.DefaultScheduledAgentManager;
import com.alibaba.cloud.ai.graph.scheduling.ScheduleConfig;
import com.alibaba.cloud.ai.graph.scheduling.ScheduleLifecycleListener;
import com.alibaba.cloud.ai.graph.scheduling.ScheduledAgentTask;

import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;

import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;

import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.jupiter.api.Assertions.fail;

/**
 * Test for ScheduledAgentTask using external TaskScheduler
 */
@EnabledIfEnvironmentVariable(named = "AI_DASHSCOPE_API_KEY", matches = ".+")
public class ScheduledAgentTaskTest {

	private ChatModel chatModel;

	@BeforeEach
	void setUp() {
		// 先创建 DashScopeApi 实例
		DashScopeApi dashScopeApi = DashScopeApi.builder().apiKey(System.getenv("AI_DASHSCOPE_API_KEY")).build();

		// 创建 DashScope ChatModel 实例
		this.chatModel = DashScopeChatModel.builder().dashScopeApi(dashScopeApi).build();
	}

	@AfterEach
	void tearDown() {
		// Clean up any active scheduled executions
		try {
			DefaultScheduledAgentManager manager = DefaultScheduledAgentManager.getInstance();
			for (String taskId : manager.getAllActiveTaskIds()) {
				manager.getTask(taskId).ifPresent(ScheduledAgentTask::stop);
			}
		}
		catch (Exception e) {
			// Ignore cleanup errors
		}
	}

	@Test
	public void testScheduleReactAgentWithCron() throws Exception {
		// Counter for tracking executions
		final AtomicInteger executionCount = new AtomicInteger(0);
		final CountDownLatch executionLatch = new CountDownLatch(2); // Wait for at least
		// 1 execution

		// Create lifecycle listener
		ScheduleLifecycleListener listener = new ScheduleLifecycleListener() {
			@Override
			public void onEvent(ScheduleEvent event, Object data) {
				if (event == ScheduleEvent.EXECUTION_COMPLETED) {
					int count = executionCount.incrementAndGet();
					System.out.println("Cron execution #" + count + " " + event + " at " + System.currentTimeMillis());
					if (data instanceof OverAllState) {
						String textContent = ((OverAllState) data).value("output")
							.filter(AssistantMessage.class::isInstance)
							.map(AssistantMessage.class::cast)
							.map(AssistantMessage::getText)
							.orElse("No content available");
						System.out.println("output:" + textContent);
					}
					System.out.println((2 - executionLatch.getCount())
							+ ": -------------------------------------------------------------------------------------\n\n");
					executionLatch.countDown();
				}
				if (event == ScheduleEvent.EXECUTION_FAILED) {
					int count = executionCount.incrementAndGet();
					System.out.println("Cron execution #" + count + " " + event + " at " + System.currentTimeMillis());
					if (data instanceof Exception) {
						fail((Exception) data);
					}
					executionLatch.countDown();
				}
			}
		};

		try {
			// Create ReactAgent
			ReactAgent agent = ReactAgent.builder()
				.name("cron_scheduled_agent")
				.model(chatModel)
				.instruction("You are a cron test agent. respond briefly.")
				.outputKey("output")
				.build();

			long startTime = System.currentTimeMillis();
			System.out.println("Cron test started at: " + startTime);

			// Schedule with cron expression - every 5 seconds
			ScheduleConfig config = ScheduleConfig.builder()
				.cronExpression("*/15 * * * * ?") // Every 5 seconds
				.inputs(Map.of("messages", List.of(new UserMessage("Is '0 */1 * * * ?' a right expression?"))))
				.addListener(listener)
				.build();

			// Schedule the agent
			ScheduledAgentTask task = agent.schedule(config);

			// Verify task is registered and started
			assertThat(task).isNotNull();
			assertThat(task.isStarted()).isTrue();

			System.out.println("Cron scheduled task with ID: " + task.getTaskId());
			System.out.println("Waiting for cron execution (max 2 MINUTES)...");

			// Wait for execution with timeout (2 MINUTES should be enough)
			boolean executionCompleted = executionLatch.await(2, TimeUnit.MINUTES);

			// Stop the task
			task.stop();

			long endTime = System.currentTimeMillis();
			System.out.println("Cron test completed at: " + endTime + ", duration: " + (endTime - startTime) + "ms");

			// Verify at least one execution occurred
			assertThat(executionCompleted).as("Expected at least 1 cron execution within 10 seconds").isTrue();

			assertThat(executionCount.get()).as("Should have at least 1 execution").isGreaterThanOrEqualTo(1);

			// Verify task is stopped
			assertThat(task.isStopped()).isTrue();

			System.out.println("✅ Cron test completed successfully - executed " + executionCount.get() + " times");

		}
		catch (Exception e) {
			fail("testScheduleReactAgentWithCron failed: " + e.getMessage());
		}
	}

}
