/*
 * Copyright 2025 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
 *
 *      http://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 org.gradle.api.internal.tasks.testing.results.serializable;

import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableSet;
import com.google.common.io.ByteStreams;
import org.gradle.api.tasks.testing.TestOutputEvent;
import org.gradle.internal.UncheckedException;
import org.gradle.internal.concurrent.CompositeStoppable;
import org.gradle.internal.io.IoConsumer;
import org.gradle.internal.serialize.Decoder;
import org.gradle.internal.serialize.Serializer;
import org.gradle.internal.serialize.kryo.KryoBackedDecoder;

import java.io.Closeable;
import java.io.EOFException;
import java.io.IOException;
import java.io.InputStream;
import java.nio.channels.Channels;
import java.nio.channels.SeekableByteChannel;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.function.LongPredicate;

/**
 * Reads test output from an output events file generated by {@link TestOutputWriter}.
 */
public final class TestOutputReader implements Closeable {
    private final Path outputEventsFile;
    private final Serializer<TestOutputEvent> testOutputEventSerializer;
    /**
     * Channels that are currently open and can be reused. We limit the pool size to avoid excessive resource usage.
     */
    private final BlockingQueue<SeekableByteChannel> channelPool = new LinkedBlockingQueue<>(64);

    TestOutputReader(Path outputEventsFile, Serializer<TestOutputEvent> testOutputEventSerializer) {
        this.outputEventsFile = outputEventsFile;
        this.testOutputEventSerializer = testOutputEventSerializer;
    }

    public boolean hasOutput(OutputEntry entry, TestOutputEvent.Destination destination) {
        switch (destination) {
            case StdOut:
                return entry.getOutputRanges().startStdout != OutputRanges.NO_OUTPUT;
            case StdErr:
                return entry.getOutputRanges().startStderr != OutputRanges.NO_OUTPUT;
            default:
                throw new IllegalArgumentException("Unknown destination: " + destination);
        }
    }

    public void useTestOutputEvents(OutputEntry entry, TestOutputEvent.Destination destination, IoConsumer<TestOutputEvent> eventConsumer) throws IOException {
        long entryStart = getStart(destination, entry);
        if (entryStart == OutputRanges.NO_OUTPUT) {
            return;
        }

        useTestOutputEventsWithInfo(destination, eventConsumer, entryStart, entry.getOutputRanges().end, fileId -> fileId == entry.getId());
    }

    public void useTestOutputEvents(Iterable<OutputEntry> entries, TestOutputEvent.Destination destination, IoConsumer<TestOutputEvent> eventConsumer) throws IOException {
        CombinedEntryInfo info = getCombinedEntryInfo(entries, destination);
        if (info.start == OutputRanges.NO_OUTPUT) {
            return;
        }
        ImmutableSet<Long> ids = info.idsBuilder.build();

        useTestOutputEventsWithInfo(destination, eventConsumer, info.start, info.end, ids::contains);
    }

    private static final class CombinedEntryInfo {
        long start = OutputRanges.NO_OUTPUT;
        long end = OutputRanges.NO_OUTPUT;
        final ImmutableSet.Builder<Long> idsBuilder = ImmutableSet.builder();
    }

    private static CombinedEntryInfo getCombinedEntryInfo(Iterable<OutputEntry> entries, TestOutputEvent.Destination destination) {
        CombinedEntryInfo info = new CombinedEntryInfo();
        for (OutputEntry entry : entries) {
            long rangesStart = getStart(destination, entry);
            if (rangesStart == OutputRanges.NO_OUTPUT) {
                continue;
            }
            if (info.start == OutputRanges.NO_OUTPUT || rangesStart < info.start) {
                info.start = rangesStart;
            }
            if (info.end == OutputRanges.NO_OUTPUT || entry.getOutputRanges().end > info.end) {
                info.end = entry.getOutputRanges().end;
            }
            info.idsBuilder.add(entry.getId());
        }
        return info;
    }

    private void useTestOutputEventsWithInfo(
        TestOutputEvent.Destination destination,
        IoConsumer<TestOutputEvent> eventConsumer,
        long start,
        long end,
        LongPredicate matchesId
    ) throws IOException {
        SeekableByteChannel channel = requestChannel();
        try {
            channel.position(start);
            InputStream stream = ByteStreams.limit(Channels.newInputStream(channel), end - start);
            Decoder decoder = new KryoBackedDecoder(stream);

            iterateEvents(decoder, matchesId, destination, eventConsumer);
        } finally {
            returnChannel(channel);
        }
    }

    private static long getStart(TestOutputEvent.Destination destination, OutputEntry entry) {
        switch (destination) {
            case StdOut:
                return entry.getOutputRanges().startStdout;
            case StdErr:
                return entry.getOutputRanges().startStderr;
            default:
                throw new IllegalArgumentException("Unknown destination: " + destination);
        }
    }

    private void iterateEvents(
        Decoder decoder,
        LongPredicate matchesId,
        TestOutputEvent.Destination destination,
        IoConsumer<TestOutputEvent> eventConsumer
    ) throws IOException {
        while (true) {
            long id;
            try {
                id = decoder.readLong();
            } catch (EOFException e) {
                // Expected EOF
                break;
            }
            TestOutputEvent event;
            try {
                event = testOutputEventSerializer.read(decoder);
            } catch (EOFException e) {
                throw new IllegalStateException("Should have reached EOF when reading the id, not in the middle of an event", e);
            } catch (Exception e) {
                Throwables.throwIfInstanceOf(e, IOException.class);
                throw UncheckedException.throwAsUncheckedException(e);
            }
            // This must be done after reading the event to ensure the decoder is positioned correctly for the next read
            if (!matchesId.test(id)) {
                continue;
            }
            if (event.getDestination() != destination) {
                continue;
            }
            eventConsumer.accept(event);
        }
    }

    private SeekableByteChannel requestChannel() {
        SeekableByteChannel open = channelPool.poll();
        if (open != null && open.isOpen()) {
            return open;
        }
        try {
            return Files.newByteChannel(outputEventsFile);
        } catch (IOException e) {
            throw new RuntimeException("Could not open output events file: " + outputEventsFile, e);
        }
    }

    private void returnChannel(SeekableByteChannel channel) {
        // Return the channel to the pool for reuse
        boolean added;
        try {
            added = channelPool.offer(channel);
        } catch (Throwable t) {
            // If we fail to return the channel to the pool, close it to avoid resource leaks
            try {
                channel.close();
            } catch (IOException e) {
                t.addSuppressed(e);
            }
            throw t;
        }
        if (!added) {
            // Pool is full, close the channel
            try {
                channel.close();
            } catch (IOException e) {
                throw UncheckedException.throwAsUncheckedException(e);
            }
        }
    }

    @Override
    public void close() throws IOException {
        CompositeStoppable.stoppable(channelPool).stop();
    }
}
