/*
 * 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.scheduling;

import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.ReentrantReadWriteLock;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.scheduling.TaskScheduler;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;

/**
 * Default manager for all scheduled agent executions. Provides a shared task scheduler
 * pool and manages all ScheduledAgentTask instances.
 *
 * @author yaohui &#064;create 2025/8/21
 */
public class DefaultScheduledAgentManager implements ScheduledAgentManager {

	private static final Logger log = LoggerFactory.getLogger(DefaultScheduledAgentManager.class);

	private static final DefaultScheduledAgentManager INSTANCE = new DefaultScheduledAgentManager();

	private final TaskScheduler taskScheduler;

	private final Map<String, ScheduledAgentTask> activeTasks = new ConcurrentHashMap<>();

	private final AtomicInteger taskIdGenerator = new AtomicInteger(1);

	private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();

	private volatile boolean shutdown = false;

	private DefaultScheduledAgentManager() {
		this.taskScheduler = createTaskScheduler();
		log.info("Default Scheduled Agent Manager initialized with shared task scheduler");
	}

	/**
	 * Get the singleton instance of the default manager
	 */
	public static DefaultScheduledAgentManager getInstance() {
		return INSTANCE;
	}

	/**
	 * Get the task scheduler
	 */
	@Override
	public TaskScheduler getTaskScheduler() {
		return taskScheduler;
	}

	/**
	 * Register a new scheduled agent execution
	 * @param task the ScheduledAgentTask to register
	 * @return the unique Task ID assigned to this Task
	 */
	@Override
	public String registerTask(ScheduledAgentTask task) {
		if (shutdown) {
			throw new IllegalStateException("Default Scheduled Agent Manager is shut down");
		}
		String taskId = String.format("agent-task-%s-%d", task.getName(), taskIdGenerator.getAndIncrement());
		lock.writeLock().lock();
		try {
			activeTasks.put(taskId, task);
			log.debug("Registered scheduled agent task: {}", taskId);
			return taskId;
		}
		finally {
			lock.writeLock().unlock();
		}
	}

	/**
	 * Unregister a scheduled agent execution
	 * @param taskId the task ID to unregister
	 * @return true if the task was found and removed, false otherwise
	 */
	@Override
	public boolean unregisterTask(String taskId) {
		lock.writeLock().lock();
		try {
			ScheduledAgentTask removed = activeTasks.remove(taskId);
			if (removed != null) {
				log.debug("Unregistered scheduled agent task: {}", taskId);
				return true;
			}
			return false;
		}
		finally {
			lock.writeLock().unlock();
		}
	}

	/**
	 * Get a scheduled agent task by ID
	 * @param taskId the task ID
	 * @return the ScheduledAgentTask if found, empty Optional otherwise
	 */
	@Override
	public Optional<ScheduledAgentTask> getTask(String taskId) {
		lock.readLock().lock();
		try {
			return Optional.ofNullable(activeTasks.get(taskId));
		}
		finally {
			lock.readLock().unlock();
		}
	}

	/**
	 * Get all active execution IDs
	 * @return a set of all active execution IDs
	 */
	@Override
	public Set<String> getAllActiveTaskIds() {
		lock.readLock().lock();
		try {
			return Set.copyOf(activeTasks.keySet());
		}
		finally {
			lock.readLock().unlock();
		}
	}

	/**
	 * Get the count of active executions
	 * @return the number of active executions
	 */
	@Override
	public int getActiveTaskCount() {
		lock.readLock().lock();
		try {
			return activeTasks.size();
		}
		finally {
			lock.readLock().unlock();
		}
	}

	/**
	 * Stop all active executions and shutdown the manager
	 */
	@Override
	public void shutdown() {
		lock.writeLock().lock();
		try {
			if (shutdown) {
				return;
			}

			log.info("Shutting down Default Scheduled Agent Manager with {} active executions", activeTasks.size());

			// Stop all active executions
			activeTasks.values().forEach(task -> {
				try {
					task.stop();
				}
				catch (Exception e) {
					log.warn("Error stopping scheduled execution during shutdown", e);
				}
			});

			activeTasks.clear();

			// Shutdown the task scheduler
			if (taskScheduler instanceof ThreadPoolTaskScheduler) {
				((ThreadPoolTaskScheduler) taskScheduler).shutdown();
			}

			shutdown = true;
			log.info("Default Scheduled Agent Manager shut down successfully");
		}
		finally {
			lock.writeLock().unlock();
		}
	}

	/**
	 * Check if the manager is shut down
	 */
	@Override
	public boolean isShutdown() {
		return shutdown;
	}

	private TaskScheduler createTaskScheduler() {
		ThreadPoolTaskScheduler scheduler = new ThreadPoolTaskScheduler();
		scheduler.setPoolSize(3);
		scheduler.setThreadNamePrefix("agent-scheduler-");
		scheduler.setThreadGroupName("DefaultAgentScheduler");
		scheduler.setWaitForTasksToCompleteOnShutdown(true);
		scheduler.setAwaitTerminationSeconds(30);
		scheduler.initialize();
		log.info("Task scheduler initialized with pool size: {}", scheduler.getPoolSize());
		return scheduler;
	}

}
