package com.easyagents.federation.sql.runtime;

import com.easyagents.federation.sql.api.FederationSqlErrorCode;
import com.easyagents.federation.sql.execute.FederationColumn;
import com.easyagents.federation.sql.execute.FederationResultCursor;
import com.easyagents.federation.sql.execute.QueryId;
import java.io.IOException;
import java.io.InputStream;
import java.io.Reader;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import org.junit.Assert;
import org.junit.Test;

/**
 * Core 托管游标的大字段流异常资源释放测试。
 */
public class ManagedFederationResultCursorTest {

    /**
     * 验证二进制流在获取后的读取失败会释放 Adapter 与 Core 资源。
     *
     * @throws Exception 测试流读取失败
     */
    @Test
    public void shouldReleaseResourcesWhenBinaryStreamReadFails() throws Exception {
        assertStreamFailure(true);
    }

    /**
     * 验证字符流在获取后的读取失败会释放 Adapter 与 Core 资源。
     *
     * @throws Exception 测试流读取失败
     */
    @Test
    public void shouldReleaseResourcesWhenCharacterStreamReadFails() throws Exception {
        assertStreamFailure(false);
    }

    /**
     * 验证关闭游标与取消并发时不会为复用的 QueryId 留下伪预取消状态。
     */
    @Test
    public void shouldRemoveCursorTrackingBeforeReleasingRegistration() {
        QueryCancellationRegistry registry = new QueryCancellationRegistry();
        QueryId queryId = new QueryId("reusable-query-id");
        QueryCancellationRegistry.QueryRegistration registration = registry.begin(queryId);
        ManagedFederationResultCursor cursor = new ManagedFederationResultCursor(
            new TestCursor(),
            registration,
            () -> Assert.assertTrue(registry.cancel(queryId))
        );

        cursor.close();

        QueryCancellationRegistry.QueryRegistration reused = registry.begin(queryId);
        reused.close();
        registry.close();
    }

    private static void assertStreamFailure(boolean binary) throws Exception {
        TestCursor delegate = new TestCursor();
        AtomicInteger resourceCloses = new AtomicInteger();
        AtomicInteger trackingCloses = new AtomicInteger();
        ManagedFederationResultCursor cursor = new ManagedFederationResultCursor(
            delegate,
            resourceCloses::incrementAndGet,
            trackingCloses::incrementAndGet
        );

        try {
            if (binary) {
                cursor.getBinaryStream(1).read();
            } else {
                cursor.getCharacterStream(1).read();
            }
            Assert.fail("LOB stream read should fail");
        } catch (IOException exception) {
            Assert.assertEquals("stream failed", exception.getMessage());
        }

        Assert.assertEquals(1, delegate.closeCalls.get());
        Assert.assertEquals(1, resourceCloses.get());
        Assert.assertEquals(1, trackingCloses.get());
        Assert.assertEquals(FederationSqlErrorCode.EXECUTION_FAILED, delegate.failure.get());
        cursor.close();
        Assert.assertEquals(1, delegate.closeCalls.get());
    }

    private static final class TestCursor
        implements FederationResultCursor, CancellationAwareFederationCursor {

        private final AtomicInteger closeCalls = new AtomicInteger();
        private final AtomicReference<FederationSqlErrorCode> failure = new AtomicReference<>();

        @Override
        public QueryId queryId() {
            return new QueryId("lob-stream-test");
        }

        @Override
        public List<FederationColumn> columns() {
            return List.of();
        }

        @Override
        public boolean next() {
            return false;
        }

        @Override
        public Object getObject(int columnIndex) {
            return null;
        }

        @Override
        public InputStream getBinaryStream(int columnIndex) {
            return new InputStream() {
                @Override
                public int read() throws IOException {
                    throw new IOException("stream failed");
                }
            };
        }

        @Override
        public Reader getCharacterStream(int columnIndex) {
            return new Reader() {
                @Override
                public int read(char[] characters, int offset, int length) throws IOException {
                    throw new IOException("stream failed");
                }

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

        @Override
        public List<Object> row() {
            return List.of();
        }

        @Override
        public void close() {
            closeCalls.incrementAndGet();
        }

        @Override
        public void markCancelled() {
        }

        @Override
        public void markTimedOut() {
        }

        @Override
        public void markFailed(FederationSqlErrorCode errorCode) {
            failure.set(errorCode);
        }
    }
}
