/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.apache.flink.connector.http.table.lookup;

import org.apache.flink.api.common.serialization.DeserializationSchema;
import org.apache.flink.api.common.serialization.SerializationSchema;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.connector.http.config.HttpConnectorConfigConstants;
import org.apache.flink.connector.http.preprocessor.HeaderPreprocessor;
import org.apache.flink.connector.http.table.lookup.querycreators.GenericGetQueryCreator;
import org.apache.flink.connector.http.table.lookup.querycreators.GenericJsonQueryCreator;
import org.apache.flink.connector.http.utils.HttpHeaderUtils;
import org.apache.flink.formats.json.JsonFormatFactory;
import org.apache.flink.table.api.DataTypes;
import org.apache.flink.table.data.GenericRowData;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.data.StringData;
import org.apache.flink.table.factories.DynamicTableFactory;
import org.apache.flink.table.types.DataType;
import org.apache.flink.util.Collector;
import org.apache.flink.util.ConfigurationException;

import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;

import java.io.IOException;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;

import static org.apache.flink.connector.http.TestHelper.assertPropertyArray;
import static org.apache.flink.connector.http.table.lookup.HttpLookupTableSourceFactory.row;
import static org.assertj.core.api.Assertions.assertThat;

/** Test for {@link JavaNetHttpPollingClient}. */
@ExtendWith(MockitoExtension.class)
public class JavaNetHttpPollingClientTest {

    @Mock private HttpClient httpClient;

    @Mock private DeserializationSchema<RowData> decoder;

    @Mock private LookupRow lookupRow;

    @Mock private DynamicTableFactory.Context dynamicTableFactoryContext;
    private HeaderPreprocessor headerPreprocessor;

    private HttpLookupConfig options;

    private static final String BASE_URL = "http://localhost.com";

    @BeforeEach
    public void setUp() {
        this.headerPreprocessor = HttpHeaderUtils.createBasicAuthorizationHeaderPreprocessor();
        this.options = HttpLookupConfig.builder().url(BASE_URL).build();
    }

    @Test
    public void shouldBuildClientWithoutHeaders() throws ConfigurationException {

        JavaNetHttpPollingClient client =
                new JavaNetHttpPollingClient(
                        httpClient,
                        decoder,
                        options,
                        new GetRequestFactory(
                                new GenericGetQueryCreator(lookupRow),
                                headerPreprocessor,
                                options));

        assertThat(((GetRequestFactory) client.getRequestFactory()).getHeadersAndValues())
                .isEmpty();
    }

    @Test
    public void shouldBuildGetClientUri() throws ConfigurationException {
        // GIVEN
        JavaNetHttpPollingClient client =
                new JavaNetHttpPollingClient(
                        httpClient,
                        decoder,
                        options,
                        new GetRequestFactory(
                                new GenericGetQueryCreator(lookupRow),
                                headerPreprocessor,
                                options));

        DataType lookupPhysicalDataType =
                row(
                        List.of(
                                DataTypes.FIELD("id", DataTypes.STRING()),
                                DataTypes.FIELD("uuid", DataTypes.STRING())));

        RowData lookupRowData =
                GenericRowData.of(StringData.fromString("1"), StringData.fromString("2"));

        LookupRow lookupRow =
                new LookupRow()
                        .addLookupEntry(
                                new RowDataSingleValueLookupSchemaEntry(
                                        "id",
                                        RowData.createFieldGetter(
                                                DataTypes.STRING().getLogicalType(), 0)))
                        .addLookupEntry(
                                new RowDataSingleValueLookupSchemaEntry(
                                        "uuid",
                                        RowData.createFieldGetter(
                                                DataTypes.STRING().getLogicalType(), 1)));
        lookupRow.setLookupPhysicalRowDataType(lookupPhysicalDataType);

        GenericGetQueryCreator queryCreator = new GenericGetQueryCreator(lookupRow);
        LookupQueryInfo lookupQueryInfo = queryCreator.createLookupQuery(lookupRowData);

        // WHEN
        URI uri = ((GetRequestFactory) client.getRequestFactory()).constructGetUri(lookupQueryInfo);

        // THEN
        assertThat(uri.toString()).isEqualTo(BASE_URL + "?id=1&uuid=2");
    }

    @Test
    public void shouldBuildBodyBasedClientUri() {
        // GIVEN
        DataType lookupPhysicalDataType =
                row(
                        List.of(
                                DataTypes.FIELD("id", DataTypes.STRING()),
                                DataTypes.FIELD("uuid", DataTypes.STRING())));

        LookupRow lookupRow =
                new LookupRow()
                        .addLookupEntry(
                                new RowDataSingleValueLookupSchemaEntry(
                                        "id",
                                        RowData.createFieldGetter(
                                                DataTypes.STRING().getLogicalType(), 0)))
                        .addLookupEntry(
                                new RowDataSingleValueLookupSchemaEntry(
                                        "uuid",
                                        RowData.createFieldGetter(
                                                DataTypes.STRING().getLogicalType(), 1)));

        lookupRow.setLookupPhysicalRowDataType(lookupPhysicalDataType);

        SerializationSchema<RowData> jsonSerializer =
                new JsonFormatFactory()
                        .createEncodingFormat(dynamicTableFactoryContext, new Configuration())
                        .createRuntimeEncoder(null, lookupPhysicalDataType);

        BodyBasedRequestFactory requestFactory =
                new BodyBasedRequestFactory(
                        "POST",
                        new GenericJsonQueryCreator(jsonSerializer),
                        HttpHeaderUtils.createBasicAuthorizationHeaderPreprocessor(),
                        HttpLookupConfig.builder().url(BASE_URL).build());

        Map<String, String> urlBodyBasedQueryParameters = new LinkedHashMap<>();
        urlBodyBasedQueryParameters.put("key1", "value1");
        urlBodyBasedQueryParameters.put("key2", "value2");

        LookupQueryInfo lookupQueryInfo =
                new LookupQueryInfo("{}", urlBodyBasedQueryParameters, null);

        // WHEN
        HttpRequest httpRequest = requestFactory.setUpRequestMethod(lookupQueryInfo).build();

        // THEN
        assertThat(httpRequest.uri().toString()).isEqualTo(BASE_URL + "?key1=value1&key2=value2");
    }

    @Test
    public void shouldBuildClientWithHeaders() throws ConfigurationException {

        // GIVEN
        Properties properties = new Properties();
        properties.setProperty("property", "val1");
        properties.setProperty("my.property", "val2");
        properties.setProperty(
                HttpConnectorConfigConstants.LOOKUP_SOURCE_HEADER_PREFIX + "Origin",
                "https://developer.mozilla.org");

        properties.setProperty(
                HttpConnectorConfigConstants.LOOKUP_SOURCE_HEADER_PREFIX + "Cache-Control",
                "no-cache, no-store, max-age=0, must-revalidate");
        properties.setProperty(
                HttpConnectorConfigConstants.LOOKUP_SOURCE_HEADER_PREFIX
                        + "Access-Control-Allow-Origin",
                "*");

        // WHEN
        HttpLookupConfig lookupConfig = HttpLookupConfig.builder().properties(properties).build();

        JavaNetHttpPollingClient client =
                new JavaNetHttpPollingClient(
                        httpClient,
                        decoder,
                        lookupConfig,
                        new GetRequestFactory(
                                new GenericGetQueryCreator(lookupRow),
                                headerPreprocessor,
                                lookupConfig));

        String[] headersAndValues =
                ((GetRequestFactory) client.getRequestFactory()).getHeadersAndValues();
        assertThat(headersAndValues).hasSize(6);

        // THEN
        // assert that we have property followed by its value.
        assertPropertyArray(headersAndValues, "Origin", "https://developer.mozilla.org");
        assertPropertyArray(
                headersAndValues,
                "Cache-Control",
                "no-cache, no-store, max-age=0, must-revalidate");
        assertPropertyArray(headersAndValues, "Access-Control-Allow-Origin", "*");
    }

    @Test
    public void deserializeSingleValueTest() throws ConfigurationException, IOException {
        // GIVEN
        DeserializationSchema<RowData> mockDecoder =
                new DeserializationSchema<RowData>() {
                    @Override
                    public RowData deserialize(byte[] message) throws IOException {
                        return null;
                    }

                    @Override
                    public void deserialize(byte[] message, Collector<RowData> out)
                            throws IOException {
                        String msg = new String(message);
                        out.collect(GenericRowData.of(StringData.fromString(msg)));
                        out.close();
                    }

                    @Override
                    public boolean isEndOfStream(RowData nextElement) {
                        return false;
                    }

                    @Override
                    public org.apache.flink.api.common.typeinfo.TypeInformation<RowData>
                            getProducedType() {
                        return null;
                    }
                };

        JavaNetHttpPollingClient client =
                new JavaNetHttpPollingClient(
                        httpClient,
                        mockDecoder,
                        options,
                        new GetRequestFactory(
                                new GenericGetQueryCreator(lookupRow),
                                headerPreprocessor,
                                options));

        // WHEN
        String testString = "Test1";
        List<RowData> result = client.deserializeSingleValue(testString.getBytes());

        // THEN
        assertThat(result).hasSize(1);
        assertThat(((StringData) result.get(0).getString(0)).toString()).isEqualTo(testString);
    }

    @Test
    public void shouldDeserializeArrayWithValidObjects() throws Exception {
        // GIVEN
        DeserializationSchema<RowData> mockDecoder =
                new DeserializationSchema<RowData>() {
                    @Override
                    public RowData deserialize(byte[] message) throws IOException {
                        return null;
                    }

                    @Override
                    public void deserialize(byte[] message, Collector<RowData> out)
                            throws IOException {
                        String msg = new String(message);
                        if (msg.contains("value1")) {
                            out.collect(GenericRowData.of(StringData.fromString("row1")));
                        } else if (msg.contains("value2")) {
                            out.collect(GenericRowData.of(StringData.fromString("row2")));
                        }
                        out.close();
                    }

                    @Override
                    public boolean isEndOfStream(RowData nextElement) {
                        return false;
                    }

                    @Override
                    public org.apache.flink.api.common.typeinfo.TypeInformation<RowData>
                            getProducedType() {
                        return null;
                    }
                };

        Properties properties = new Properties();
        properties.setProperty(HttpConnectorConfigConstants.RESULT_TYPE, "array");

        HttpLookupConfig lookupConfig =
                HttpLookupConfig.builder().url(BASE_URL).properties(properties).build();

        JavaNetHttpPollingClient client =
                new JavaNetHttpPollingClient(
                        httpClient,
                        mockDecoder,
                        lookupConfig,
                        new GetRequestFactory(
                                new GenericGetQueryCreator(lookupRow),
                                headerPreprocessor,
                                lookupConfig));

        // WHEN
        String jsonArray = "[{\"key\":\"value1\"},{\"key\":\"value2\"}]";
        List<RowData> result = client.deserializeArray(jsonArray.getBytes());

        // THEN
        assertThat(result).isNotNull();
        assertThat(result).hasSize(2);
    }

    @Test
    public void shouldHandleNullNodesInArray() throws Exception {
        // GIVEN
        DeserializationSchema<RowData> mockDecoder =
                new DeserializationSchema<RowData>() {
                    @Override
                    public RowData deserialize(byte[] message) throws IOException {
                        return null;
                    }

                    @Override
                    public void deserialize(byte[] message, Collector<RowData> out)
                            throws IOException {
                        out.collect(GenericRowData.of(StringData.fromString("valid")));
                        out.close();
                    }

                    @Override
                    public boolean isEndOfStream(RowData nextElement) {
                        return false;
                    }

                    @Override
                    public org.apache.flink.api.common.typeinfo.TypeInformation<RowData>
                            getProducedType() {
                        return null;
                    }
                };

        Properties properties = new Properties();
        properties.setProperty(HttpConnectorConfigConstants.RESULT_TYPE, "array");

        HttpLookupConfig lookupConfig =
                HttpLookupConfig.builder().url(BASE_URL).properties(properties).build();

        JavaNetHttpPollingClient client =
                new JavaNetHttpPollingClient(
                        httpClient,
                        mockDecoder,
                        lookupConfig,
                        new GetRequestFactory(
                                new GenericGetQueryCreator(lookupRow),
                                headerPreprocessor,
                                lookupConfig));

        // WHEN
        String jsonArray = "[{\"key\":\"value1\"},null,{\"key\":\"value2\"}]";
        List<RowData> result = client.deserializeArray(jsonArray.getBytes());

        // THEN - null nodes should be skipped
        assertThat(result).isNotNull();
        assertThat(result).hasSize(2);
    }

    @Test
    public void shouldHandleEmptyDeserializationInArray() throws Exception {
        // GIVEN
        DeserializationSchema<RowData> mockDecoder =
                new DeserializationSchema<RowData>() {
                    @Override
                    public RowData deserialize(byte[] message) throws IOException {
                        return null;
                    }

                    @Override
                    public void deserialize(byte[] message, Collector<RowData> out)
                            throws IOException {
                        String msg = new String(message);
                        // Only collect for specific messages, return empty for others
                        if (msg.contains("\"status\":\"valid\"")) {
                            out.collect(GenericRowData.of(StringData.fromString("data")));
                        }
                        // Don't collect anything for other messages
                        out.close();
                    }

                    @Override
                    public boolean isEndOfStream(RowData nextElement) {
                        return false;
                    }

                    @Override
                    public org.apache.flink.api.common.typeinfo.TypeInformation<RowData>
                            getProducedType() {
                        return null;
                    }
                };

        Properties properties = new Properties();
        properties.setProperty(HttpConnectorConfigConstants.RESULT_TYPE, "array");

        HttpLookupConfig lookupConfig =
                HttpLookupConfig.builder().url(BASE_URL).properties(properties).build();

        JavaNetHttpPollingClient client =
                new JavaNetHttpPollingClient(
                        httpClient,
                        mockDecoder,
                        lookupConfig,
                        new GetRequestFactory(
                                new GenericGetQueryCreator(lookupRow),
                                headerPreprocessor,
                                lookupConfig));

        // WHEN
        String jsonArray =
                "[{\"status\":\"invalid\"},{\"status\":\"valid\"},{\"status\":\"invalid\"}]";
        List<RowData> result = client.deserializeArray(jsonArray.getBytes());

        // THEN - only valid deserialization should be included
        assertThat(result).isNotNull();
        assertThat(result).hasSize(1);
    }
}
