/*
 * 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.time.Duration;
import java.time.Instant;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ScheduledFuture;

import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.OverAllStateBuilder;
import com.alibaba.cloud.ai.graph.RunnableConfig;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.springframework.scheduling.TaskScheduler;
import org.springframework.scheduling.support.CronTrigger;

/**
 * ScheduledGraphExecution
 *
 * @author yaohui &#064;create 2025/8/20 15:04
 **/
public class ScheduledAgentTask {

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

	private final TaskScheduler taskScheduler;

	private final CompiledGraph graph;

	private final ScheduleConfig config;

	private volatile ScheduledFuture<?> scheduledFuture;

	private volatile boolean started = false;

	private volatile boolean stopped = false;

	private final String taskId;

	public ScheduledAgentTask(CompiledGraph graph, ScheduleConfig config) {
		this.graph = graph;
		this.config = config;
		ScheduledAgentManager scheduledAgentManager = ScheduledAgentManagerFactory.getInstance().getManager();
		this.taskScheduler = scheduledAgentManager.getTaskScheduler();
		// Register with the active manager
		this.taskId = scheduledAgentManager.registerTask(this);
		log.debug("Created ScheduledAgentTask with ID: {}", taskId);
	}

	/**
	 * Start the scheduled execution
	 */
	public ScheduledAgentTask start() {
		if (started) {
			throw new IllegalStateException("Schedule already started");
		}

		switch (config.getMode()) {
			case CRON:
				scheduledFuture = taskScheduler.schedule(this::executeGraph,
						new CronTrigger(config.getCronExpression()));
				break;
			case FIXED_DELAY:
				scheduledFuture = taskScheduler.scheduleWithFixedDelay(this::executeGraph,
						Instant.now().plusMillis(config.getInitialDelay()), Duration.ofMillis(config.getFixedDelay()));
				break;
			case FIXED_RATE:
				scheduledFuture = taskScheduler.scheduleAtFixedRate(this::executeGraph,
						Instant.now().plusMillis(config.getInitialDelay()), Duration.ofMillis(config.getFixedRate()));
				break;
			case ONE_TIME:
				scheduledFuture = taskScheduler.schedule(this::executeGraph,
						Instant.now().plusMillis(config.getInitialDelay()));
				break;
			case TRIGGER:
				scheduledFuture = taskScheduler.schedule(this::executeGraph, config.getTrigger());
				break;
			default:
				throw new IllegalStateException("Unexpected value: " + config.getMode());
		}

		started = true;
		notifyListeners(ScheduleLifecycleListener.ScheduleEvent.STARTED);
		return this;
	}

	/**
	 * Stop the scheduled execution
	 */
	public void stop() {
		if (scheduledFuture != null && !scheduledFuture.isCancelled()) {
			scheduledFuture.cancel(false);
		}
		stopped = true;

		// Unregister from active manager
		ScheduledAgentManagerFactory.getInstance().getManager().unregisterTask(taskId);
		log.debug("Stopped and unregistered ScheduledAgentTask with ID: {}", taskId);

		notifyListeners(ScheduleLifecycleListener.ScheduleEvent.STOPPED);
	}

	/**
	 * Execute the graph
	 * @param runnableConfig
	 * @param inputs
	 */
	public void execute(RunnableConfig runnableConfig, Map<String, Object> inputs) {
		try {
			notifyListeners(ScheduleLifecycleListener.ScheduleEvent.EXECUTION_STARTED);
			OverAllState initialState = createInitialState(inputs);
			if (runnableConfig == null) {
				String threadId = String.format("%s-%d", taskId, System.currentTimeMillis());
				runnableConfig = RunnableConfig.builder().threadId(threadId).build();
			}
			Optional<OverAllState> result = graph.invoke(initialState, runnableConfig);
			notifyListeners(ScheduleLifecycleListener.ScheduleEvent.EXECUTION_COMPLETED, result.orElse(null));
		}
		catch (Exception e) {
			log.error("Graph execution failed", e);
			notifyListeners(ScheduleLifecycleListener.ScheduleEvent.EXECUTION_FAILED, e);
		}
	}

	/**
	 * Execute the graph with retry logic
	 */
	private void executeGraph() {
		int attempt = 0;
		while (attempt <= config.getMaxRetries()) {
			try {
				execute(config.getRunnableConfig(), config.getInputs());
				return;
			}
			catch (Exception e) {
				log.warn("Graph execution failed (attempt {}): {}", attempt + 1, e.getMessage());
				if (attempt < config.getMaxRetries() && config.getRetryPredicate().apply(e)) {
					attempt++;
					try {
						Thread.sleep(config.getRetryDelay().toMillis());
					}
					catch (InterruptedException ie) {
						Thread.currentThread().interrupt();
						break;
					}
				}
				else {
					break;
				}
			}
		}
	}

	private OverAllState createInitialState(Map<String, Object> inputs) {
		return OverAllStateBuilder.builder()
			.withKeyStrategies(graph.getKeyStrategyMap())
			.withData(inputs)
			.build();
	}

	private void notifyListeners(ScheduleLifecycleListener.ScheduleEvent event, Object data) {
		config.getListeners().forEach(listener -> {
			try {
				listener.onEvent(event, data);
			}
			catch (Exception e) {
				log.error("Error in schedule listener", e);
			}
		});
	}

	private void notifyListeners(ScheduleLifecycleListener.ScheduleEvent event) {
		notifyListeners(event, null);
	}

	// Getters for monitoring
	public boolean isStarted() {
		return started;
	}

	public boolean isStopped() {
		return stopped;
	}

	/**
	 * Get the unique task ID assigned by the global manager
	 */
	public String getTaskId() {
		return taskId;
	}

	/**
	 * Get the name of the agent associated with this task
	 */
	public String getName() {
		return graph.stateGraph.getName();
	}

}
