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

import com.alibaba.cloud.ai.graph.RunnableConfig;
import com.alibaba.cloud.ai.graph.checkpoint.BaseCheckpointSaver;
import com.alibaba.cloud.ai.graph.checkpoint.Checkpoint;
import com.alibaba.cloud.ai.graph.utils.TryFunction;

import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.locks.ReentrantLock;
import java.util.stream.IntStream;

import static java.lang.String.format;

public class MemorySaver implements BaseCheckpointSaver {

	final Map<String, LinkedList<Checkpoint>> _checkpointsByThread = new HashMap<>();
	private final ReentrantLock _lock = new ReentrantLock();

	/**
	 * Protected constructor for MemorySaver.
	 * Use {@link #builder()} to create instances.
	 */
	public MemorySaver() {
	}

	/**
	 * Creates a new builder for MemorySaver.
	 * @return a new Builder instance
	 */
	public static Builder builder() {
		return new Builder();
	}

	protected LinkedList<Checkpoint> loadedCheckpoints(RunnableConfig config, LinkedList<Checkpoint> checkpoints) throws Exception {
		return checkpoints;
	}

	protected void insertedCheckpoint(RunnableConfig config, LinkedList<Checkpoint> checkpoints, Checkpoint checkpoint) throws Exception {
	}

	protected void updatedCheckpoint(RunnableConfig config, LinkedList<Checkpoint> checkpoints, Checkpoint checkpoint) throws Exception {
	}

	protected void releasedCheckpoints(RunnableConfig config, LinkedList<Checkpoint> checkpoints, Tag releaseTag) throws Exception {
	}

	protected final <T> T loadOrInitCheckpoints(RunnableConfig config,
			TryFunction<LinkedList<Checkpoint>, T, Exception> transformer) throws Exception {
		_lock.lock();
		try {
			var threadId = config.threadId().orElse(THREAD_ID_DEFAULT);
			return transformer.tryApply(loadedCheckpoints(config, _checkpointsByThread.computeIfAbsent(threadId, k -> new LinkedList<>())));

		}
		finally {
			_lock.unlock();
		}
	}

	protected final Collection<Checkpoint> remove(String threadId) {
		return _checkpointsByThread.remove(Objects.requireNonNull(threadId));
	}

	@Override
	public final Collection<Checkpoint> list(RunnableConfig config) {
		try {
			return loadOrInitCheckpoints(config, Collections::unmodifiableCollection);
		}
		catch (Exception e) {
			throw new RuntimeException(e);
		}
	}

	@Override
	public final Optional<Checkpoint> get(RunnableConfig config) {

		try {
			return loadOrInitCheckpoints(config, checkpoints -> {
				if (config.checkPointId().isPresent()) {
					return config.checkPointId()
							.flatMap(id -> checkpoints.stream()
									.filter(checkpoint -> checkpoint.getId().equals(id))
									.findFirst());
				}
				return getLast(checkpoints, config);

			});
		}
		catch (Exception e) {
			throw new RuntimeException(e);
		}
	}

	@Override
	public final RunnableConfig put(RunnableConfig config, Checkpoint checkpoint) throws Exception {

		return loadOrInitCheckpoints(config, checkpoints -> {

			if (config.checkPointId().isPresent()) { // Replace Checkpoint
				String checkPointId = config.checkPointId().get();
				int index = IntStream.range(0, checkpoints.size())
						.filter(i -> checkpoints.get(i).getId().equals(checkPointId))
						.findFirst()
						.orElseThrow(() -> (new NoSuchElementException(format("Checkpoint with id %s not found!", checkPointId))));
				checkpoints.set(index, checkpoint);
				retainLatestCheckpoints(checkpoints, config);
				updatedCheckpoint(config, checkpoints, checkpoint);
				return config;
			}

			checkpoints.push(checkpoint); // Add Checkpoint
			retainLatestCheckpoints(checkpoints, config);
			insertedCheckpoint(config, checkpoints, checkpoint);

			return RunnableConfig.builder(config)
					.checkPointId(checkpoint.getId())
					.build();

		});
	}

	@Override
	public final Tag release(RunnableConfig config) throws Exception {

		return loadOrInitCheckpoints(config, checkpoints -> {

			var threadId = config.threadId().orElse(THREAD_ID_DEFAULT);

			var tag = new Tag(threadId, remove(threadId));

			releasedCheckpoints(config, checkpoints, tag);

			return tag;
		});
	}

	/**
	 * Builder class for MemorySaver.
	 */
	public static class Builder {
		/**
		 * Builds a new MemorySaver instance.
		 * @return a new MemorySaver instance
		 */
		public MemorySaver build() {
			return new MemorySaver();
		}
	}
}
