package cn.yunzhixue.ability.center.examsprint.infrastructure.report.pdf;

import cn.yunzhixue.ability.center.examsprint.domain.report.ExamSprintReportPdfGenerator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.DisposableBean;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

public class PooledPlaywrightExamSprintReportPdfGenerator implements ExamSprintReportPdfGenerator, DisposableBean, AutoCloseable {

    private static final Logger LOGGER = LoggerFactory.getLogger(PooledPlaywrightExamSprintReportPdfGenerator.class);

    private final int poolSize;
    private final ArrayBlockingQueue<PlaywrightPdfWorker> pool;
    private final Set<PlaywrightPdfWorker> workers = ConcurrentHashMap.newKeySet();
    private final PlaywrightPdfWorkerFactory workerFactory;
    private final Duration borrowTimeout;
    private final AtomicBoolean closed = new AtomicBoolean(false);

    public PooledPlaywrightExamSprintReportPdfGenerator(int poolSize,
                                                        Duration borrowTimeout,
                                                        PlaywrightPdfWorkerFactory workerFactory) {
        if (poolSize < 1) {
            throw new IllegalArgumentException("poolSize must be at least 1");
        }
        this.poolSize = poolSize;
        this.borrowTimeout = validateBorrowTimeout(borrowTimeout);
        this.workerFactory = Objects.requireNonNull(workerFactory, "workerFactory");
        this.pool = new ArrayBlockingQueue<>(poolSize);

        List<PlaywrightPdfWorker> createdWorkers = new ArrayList<>(poolSize);
        try {
            for (int i = 0; i < poolSize; i++) {
                PlaywrightPdfWorker worker = Objects.requireNonNull(this.workerFactory.create(), "worker");
                workers.add(worker);
                pool.add(worker);
                createdWorkers.add(worker);
            }
            PoolState state = poolState();
            LOGGER.info(
                    "exam_sprint_report_pdf_pool_initialized poolSize={} availableWorkers={} activeWorkers={} managedWorkers={} borrowTimeoutMs={}",
                    state.poolSize(),
                    state.availableWorkers(),
                    state.activeWorkers(),
                    state.managedWorkers(),
                    this.borrowTimeout.toMillis());
        } catch (RuntimeException | Error failure) {
            closeWorkers(createdWorkers);
            throw failure;
        }
    }

    @Override
    public byte[] generate(String htmlContent) {
        Objects.requireNonNull(htmlContent, "htmlContent");
        ensureOpen();

        PlaywrightPdfWorker worker = borrowWorker();
        boolean success = false;
        long workerStartedNanos = System.nanoTime();
        try {
            byte[] pdfBytes = worker.generate(htmlContent);
            success = true;
            PoolState state = poolState();
            LOGGER.info(
                    "exam_sprint_report_pdf_worker_generation_completed poolSize={} availableWorkers={} activeWorkers={} managedWorkers={} workerRenderMs={} pdfByteLength={}",
                    state.poolSize(),
                    state.availableWorkers(),
                    state.activeWorkers(),
                    state.managedWorkers(),
                    elapsedMillis(workerStartedNanos),
                    pdfBytes.length);
            return pdfBytes;
        } catch (RuntimeException | Error failure) {
            PoolState state = poolState();
            LOGGER.warn(
                    "exam_sprint_report_pdf_worker_generation_failed poolSize={} availableWorkers={} activeWorkers={} managedWorkers={} workerRenderMs={} exceptionType={}",
                    state.poolSize(),
                    state.availableWorkers(),
                    state.activeWorkers(),
                    state.managedWorkers(),
                    elapsedMillis(workerStartedNanos),
                    exceptionType(failure));
            handleFailedWorker(worker, failure);
            throw failure;
        } finally {
            if (success) {
                returnWorker(worker);
            }
        }
    }

    @Override
    public void destroy() {
        close();
    }

    @Override
    public void close() {
        if (!closed.compareAndSet(false, true)) {
            return;
        }
        List<PlaywrightPdfWorker> workersToClose = new ArrayList<>(workers);
        pool.clear();
        closeWorkers(workersToClose);
        PoolState state = poolState();
        LOGGER.info(
                "exam_sprint_report_pdf_pool_closed poolSize={} availableWorkers={} activeWorkers={} managedWorkers={} closedWorkers={} remainingWorkers={}",
                state.poolSize(),
                state.availableWorkers(),
                state.activeWorkers(),
                state.managedWorkers(),
                workersToClose.size(),
                state.managedWorkers());
    }

    private PlaywrightPdfWorker borrowWorker() {
        long borrowStartedNanos = System.nanoTime();
        long timeoutNanos = borrowTimeout.toNanos();
        long deadlineNanos = System.nanoTime() + timeoutNanos;
        long pollIntervalNanos = TimeUnit.MILLISECONDS.toNanos(50);

        try {
            while (true) {
                ensureOpen();
                long remainingNanos = deadlineNanos - System.nanoTime();
                if (remainingNanos <= 0) {
                    long borrowWaitMs = elapsedMillis(borrowStartedNanos);
                    PoolState state = poolState();
                    LOGGER.warn(
                            "exam_sprint_report_pdf_worker_borrow_timeout poolSize={} availableWorkers={} activeWorkers={} managedWorkers={} borrowWaitMs={} borrowTimeoutMs={} exceptionType={}",
                            state.poolSize(),
                            state.availableWorkers(),
                            state.activeWorkers(),
                            state.managedWorkers(),
                            borrowWaitMs,
                            borrowTimeout.toMillis(),
                            IllegalStateException.class.getSimpleName());
                    throw new IllegalStateException("Timed out waiting for PDF worker after " + borrowTimeout.toMillis() + " ms");
                }

                PlaywrightPdfWorker worker = pool.poll(Math.min(remainingNanos, pollIntervalNanos), TimeUnit.NANOSECONDS);
                if (worker != null) {
                    ensureOpen();
                    PoolState state = poolState();
                    LOGGER.info(
                            "exam_sprint_report_pdf_worker_borrowed poolSize={} availableWorkers={} activeWorkers={} managedWorkers={} borrowWaitMs={}",
                            state.poolSize(),
                            state.availableWorkers(),
                            state.activeWorkers(),
                            state.managedWorkers(),
                            elapsedMillis(borrowStartedNanos));
                    return worker;
                }
            }
        } catch (InterruptedException exception) {
            Thread.currentThread().interrupt();
            if (closed.get()) {
                throw new IllegalStateException("PDF generator pool is closed", exception);
            }
            throw new IllegalStateException("Timed out waiting for PDF worker after " + borrowTimeout.toMillis() + " ms", exception);
        }
    }

    private void returnWorker(PlaywrightPdfWorker worker) {
        if (closed.get()) {
            closeManagedWorker(worker);
            return;
        }
        if (!pool.offer(worker)) {
            closeManagedWorker(worker);
            PoolState state = poolState();
            LOGGER.warn(
                    "exam_sprint_report_pdf_worker_return_failed poolSize={} availableWorkers={} activeWorkers={} managedWorkers={} reason=pool_full",
                    state.poolSize(),
                    state.availableWorkers(),
                    state.activeWorkers(),
                    state.managedWorkers());
        } else if (closed.get() && pool.remove(worker)) {
            closeManagedWorker(worker);
        } else {
            PoolState state = poolState();
            LOGGER.info(
                    "exam_sprint_report_pdf_worker_returned poolSize={} availableWorkers={} activeWorkers={} managedWorkers={}",
                    state.poolSize(),
                    state.availableWorkers(),
                    state.activeWorkers(),
                    state.managedWorkers());
        }
    }

    private void handleFailedWorker(PlaywrightPdfWorker failedWorker, Throwable generateFailure) {
        closeManagedWorker(failedWorker);
        if (closed.get()) {
            return;
        }

        try {
            PlaywrightPdfWorker replacementWorker = Objects.requireNonNull(workerFactory.create(), "worker");
            workers.add(replacementWorker);
            if (!closed.get()) {
                if (!pool.offer(replacementWorker)) {
                    closeManagedWorker(replacementWorker);
                    PoolState state = poolState();
                    LOGGER.warn(
                            "exam_sprint_report_pdf_worker_replacement_failed poolSize={} availableWorkers={} activeWorkers={} managedWorkers={} exceptionType={} reason=pool_full",
                            state.poolSize(),
                            state.availableWorkers(),
                            state.activeWorkers(),
                            state.managedWorkers(),
                            "None");
                } else if (closed.get() && pool.remove(replacementWorker)) {
                    closeManagedWorker(replacementWorker);
                } else {
                    PoolState state = poolState();
                    LOGGER.info(
                            "exam_sprint_report_pdf_worker_replaced poolSize={} availableWorkers={} activeWorkers={} managedWorkers={}",
                            state.poolSize(),
                            state.availableWorkers(),
                            state.activeWorkers(),
                            state.managedWorkers());
                }
            } else {
                closeManagedWorker(replacementWorker);
            }
        } catch (RuntimeException | Error replacementFailure) {
            PoolState state = poolState();
            LOGGER.warn(
                    "exam_sprint_report_pdf_worker_replacement_failed poolSize={} availableWorkers={} activeWorkers={} managedWorkers={} exceptionType={}",
                    state.poolSize(),
                    state.availableWorkers(),
                    state.activeWorkers(),
                    state.managedWorkers(),
                    exceptionType(replacementFailure),
                    replacementFailure);
        }
    }

    private void ensureOpen() {
        if (closed.get()) {
            throw new IllegalStateException("PDF generator pool is closed");
        }
    }

    private Duration validateBorrowTimeout(Duration borrowTimeout) {
        Duration timeout = Objects.requireNonNull(borrowTimeout, "borrowTimeout");
        if (timeout.isZero() || timeout.isNegative()) {
            throw new IllegalArgumentException("borrowTimeout must be positive");
        }
        return timeout;
    }

    private void closeWorkers(List<PlaywrightPdfWorker> workers) {
        for (PlaywrightPdfWorker worker : workers) {
            closeManagedWorker(worker);
        }
    }

    private void closeManagedWorker(PlaywrightPdfWorker worker) {
        if (workers.remove(worker)) {
            closeQuietly(worker);
        }
    }

    private void closeQuietly(PlaywrightPdfWorker worker) {
        try {
            worker.close();
        } catch (RuntimeException | Error exception) {
            PoolState state = poolState();
            LOGGER.warn(
                    "exam_sprint_report_pdf_worker_close_failed poolSize={} availableWorkers={} activeWorkers={} managedWorkers={} exceptionType={}",
                    state.poolSize(),
                    state.availableWorkers(),
                    state.activeWorkers(),
                    state.managedWorkers(),
                    exceptionType(exception),
                    exception);
        }
    }

    private long elapsedMillis(long startedNanos) {
        return TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startedNanos);
    }

    private String exceptionType(Throwable throwable) {
        return throwable.getClass().getSimpleName();
    }

    private PoolState poolState() {
        int availableWorkers = pool.size();
        int managedWorkers = workers.size();
        return new PoolState(poolSize, availableWorkers, managedWorkers, Math.max(0, managedWorkers - availableWorkers));
    }

    private record PoolState(int poolSize, int availableWorkers, int managedWorkers, int activeWorkers) {
    }
}
