/*
 * 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.studio.core.utils.concurrent;

import com.alibaba.cloud.ai.studio.runtime.domain.RequestContext;
import com.alibaba.cloud.ai.studio.core.context.RequestContextHolder;
import org.jetbrains.annotations.NotNull;

import java.util.List;
import java.util.Collection;
import java.util.concurrent.*;

/**
 * A wrapper for ExecutorService that propagates and cleans up RequestContext across
 * thread boundaries. This ensures that request context information is properly maintained
 * in asynchronous operations.
 */
public class RequestContextThreadPoolWrapper implements ExecutorService {

	/** The underlying executor service being wrapped */
	private final ExecutorService delegate;

	/**
	 * Creates a new wrapper for the given executor service
	 * @param delegate the executor service to wrap
	 */
	public RequestContextThreadPoolWrapper(ExecutorService delegate) {
		this.delegate = delegate;
	}

	/**
	 * Executes the given command with request context propagation
	 */
	@Override
	public void execute(@NotNull Runnable command) {
		RequestContext context = RequestContextHolder.getRequestContext();
		delegate.execute(() -> {
			try {
				if (context != null) {
					RequestContextHolder.setRequestContext(context);
				}
				command.run();
			}
			finally {
				RequestContextHolder.clearRequestContext();
			}
		});
	}

	/**
	 * Submits a callable task with request context propagation
	 */
	@NotNull
	@Override
	public <T> Future<T> submit(@NotNull Callable<T> task) {
		RequestContext context = RequestContextHolder.getRequestContext();
		return delegate.submit(() -> {
			try {
				if (context != null) {
					RequestContextHolder.setRequestContext(context);
				}
				return task.call();
			}
			finally {
				RequestContextHolder.clearRequestContext();
			}
		});
	}

	/**
	 * Submits a runnable task with request context propagation and returns the given
	 * result
	 */
	@NotNull
	@Override
	public <T> Future<T> submit(@NotNull Runnable task, T result) {
		RequestContext context = RequestContextHolder.getRequestContext();
		return delegate.submit(() -> {
			try {
				if (context != null) {
					RequestContextHolder.setRequestContext(context);
				}
				task.run();
			}
			finally {
				RequestContextHolder.clearRequestContext();
			}
		}, result);
	}

	/**
	 * Submits a runnable task with request context propagation
	 */
	@NotNull
	@Override
	public Future<?> submit(@NotNull Runnable task) {
		RequestContext context = RequestContextHolder.getRequestContext();
		return delegate.submit(() -> {
			try {
				if (context != null) {
					RequestContextHolder.setRequestContext(context);
				}
				task.run();
			}
			finally {
				RequestContextHolder.clearRequestContext();
			}
		});
	}

	/**
	 * Invokes all tasks with request context propagation
	 */
	@NotNull
	@Override
	public <T> List<Future<T>> invokeAll(@NotNull Collection<? extends Callable<T>> tasks) throws InterruptedException {
		RequestContext context = RequestContextHolder.getRequestContext();
		return delegate.invokeAll(wrapCallables(tasks, context));
	}

	/**
	 * Invokes all tasks with request context propagation and timeout
	 */
	@NotNull
	@Override
	public <T> List<Future<T>> invokeAll(@NotNull Collection<? extends Callable<T>> tasks, long timeout,
			@NotNull TimeUnit unit) throws InterruptedException {
		RequestContext context = RequestContextHolder.getRequestContext();
		return delegate.invokeAll(wrapCallables(tasks, context), timeout, unit);
	}

	/**
	 * Invokes any task with request context propagation
	 */
	@NotNull
	@Override
	public <T> T invokeAny(@NotNull Collection<? extends Callable<T>> tasks)
			throws InterruptedException, ExecutionException {
		RequestContext context = RequestContextHolder.getRequestContext();
		return delegate.invokeAny(wrapCallables(tasks, context));
	}

	/**
	 * Invokes any task with request context propagation and timeout
	 */
	@Override
	public <T> T invokeAny(@NotNull Collection<? extends Callable<T>> tasks, long timeout, @NotNull TimeUnit unit)
			throws InterruptedException, ExecutionException, TimeoutException {
		RequestContext context = RequestContextHolder.getRequestContext();
		return delegate.invokeAny(wrapCallables(tasks, context), timeout, unit);
	}

	/**
	 * Wraps a collection of callables with request context propagation
	 */
	private <T> Collection<? extends Callable<T>> wrapCallables(Collection<? extends Callable<T>> tasks,
			RequestContext context) {
		return tasks.stream().map(task -> (Callable<T>) () -> {
			try {
				if (context != null) {
					RequestContextHolder.setRequestContext(context);
				}
				return task.call();
			}
			finally {
				RequestContextHolder.clearRequestContext();
			}
		}).collect(java.util.stream.Collectors.toList());
	}

	/**
	 * Initiates an orderly shutdown of the underlying executor
	 */
	@Override
	public void shutdown() {
		delegate.shutdown();
	}

	/**
	 * Attempts to stop all actively executing tasks
	 */
	@NotNull
	@Override
	public List<Runnable> shutdownNow() {
		return delegate.shutdownNow();
	}

	/**
	 * Returns true if the executor has been shut down
	 */
	@Override
	public boolean isShutdown() {
		return delegate.isShutdown();
	}

	/**
	 * Returns true if all tasks have completed following shutdown
	 */
	@Override
	public boolean isTerminated() {
		return delegate.isTerminated();
	}

	/**
	 * Blocks until all tasks have completed execution after a shutdown request
	 */
	@Override
	public boolean awaitTermination(long timeout, @NotNull TimeUnit unit) throws InterruptedException {
		return delegate.awaitTermination(timeout, unit);
	}

	/**
	 * Returns the underlying ThreadPoolExecutor if available
	 * @return the underlying ThreadPoolExecutor, or null if the delegate is not a
	 * ThreadPoolExecutor
	 */
	public ThreadPoolExecutor getThreadPoolExecutor() {
		if (delegate instanceof ThreadPoolExecutor) {
			return (ThreadPoolExecutor) delegate;
		}
		return null;
	}

}
