/*
 * Copyright Red Hat, Inc. and/or its affiliates
 * and other contributors as indicated by the @author tags and
 * the COPYRIGHT.txt file distributed with this work.
 *
 * 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.teiid.translator.odata4;

import org.mockito.ArgumentCaptor;
import org.mockito.Mockito;
import org.teiid.cdk.api.TranslationUtility;
import org.teiid.core.types.InputStreamFactory;
import org.teiid.core.util.ObjectConverterUtil;
import org.teiid.language.Call;
import org.teiid.language.Command;
import org.teiid.language.QueryExpression;
import org.teiid.metadata.MetadataFactory;
import org.teiid.translator.*;

import javax.activation.DataSource;
import javax.xml.ws.Dispatch;
import javax.xml.ws.Service;
import javax.xml.ws.handler.MessageContext;
import javax.xml.ws.http.HTTPBinding;
import java.io.*;
import java.net.URLDecoder;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static org.junit.jupiter.api.Assertions.assertEquals;

/**
 * @author Muhammad Imran Ansari
 * @date 2024.11.01
 */

public class TestBaseODataMockExecution {

    protected ResultSetExecution helpQueryExecute(MetadataFactory mf, String query,
                                                  final String resultJson, String expectedURL) throws Exception {
        return helpQueryExecute(mf, query, resultJson, expectedURL, 200);
    }

    protected ResultSetExecution helpQueryExecute(MetadataFactory mf, String query,
                                                  final String resultJson, String expectedURL, int responseCode)
            throws Exception {

        ODataExecutionFactory translator = new ODataExecutionFactory();
        translator.start();

        TranslationUtility utility = new TranslationUtility(
                TestODataMetadataProcessor.getTransformationMetadata(mf, translator));

        Command cmd = utility.parseCommand(query);
        ExecutionContext context = Mockito.mock(ExecutionContext.class);
        WSConnection connection = Mockito.mock(WSConnection.class);

        Map<String, Object> headers = new HashMap<String, Object>();
        headers.put(MessageContext.HTTP_REQUEST_HEADERS, new HashMap<String, List<String>>());
        headers.put(WSConnection.RESPONSE_CODE, responseCode);

        Dispatch<DataSource> dispatch = Mockito.mock(Dispatch.class);
        Mockito.when(dispatch.getRequestContext()).thenReturn(headers);
        Mockito.when(dispatch.getResponseContext()).thenReturn(headers);

        Mockito.when(connection.createDispatch(Mockito.eq(HTTPBinding.HTTP_BINDING), Mockito.anyString(),
                Mockito.eq(DataSource.class), Mockito.eq(Service.Mode.MESSAGE), Mockito.anyBoolean())).thenReturn(dispatch);

        DataSource ds = new DataSource() {
            @Override
            public OutputStream getOutputStream() throws IOException {
                return new ByteArrayOutputStream();
            }

            @Override
            public String getName() {
                return "result";
            }

            @Override
            public InputStream getInputStream() throws IOException {
                ByteArrayInputStream in = new ByteArrayInputStream(resultJson.getBytes());
                return in;
            }

            @Override
            public String getContentType() {
                return "application/xml";
            }
        };
        Mockito.when(dispatch.invoke(Mockito.any())).thenReturn(ds);

        ResultSetExecution execution = translator
                .createResultSetExecution((QueryExpression) cmd, context,
                        utility.createRuntimeMetadata(), connection);
        execution.execute();

        ArgumentCaptor<String> endpoint = ArgumentCaptor.forClass(String.class);
        ArgumentCaptor<String> binding = ArgumentCaptor.forClass(String.class);

        Mockito.verify(connection).createDispatch(binding.capture(),
                endpoint.capture(), Mockito.eq(DataSource.class),
                Mockito.eq(Service.Mode.MESSAGE), Mockito.anyBoolean());
        assertEquals(expectedURL, URLDecoder.decode(endpoint.getValue(), "utf-8"));
        return execution;
    }

    protected ProcedureExecution helpProcedureExecute(MetadataFactory mf, String query,
                                                      final String resultJson, String expectedURL, int responseCode)
            throws Exception {
        return helpProcedureExecute(mf, query, resultJson, expectedURL, responseCode, true);
    }

    protected ProcedureExecution helpProcedureExecute(MetadataFactory mf, String query,
                                                      final String resultJson, String expectedURL, int responseCode, boolean decode)
            throws Exception {

        ODataExecutionFactory translator = new ODataExecutionFactory();
        translator.start();

        TranslationUtility utility = new TranslationUtility(
                TestODataMetadataProcessor.getTransformationMetadata(mf, translator));

        Command cmd = utility.parseCommand(query);
        ExecutionContext context = Mockito.mock(ExecutionContext.class);
        WSConnection connection = Mockito.mock(WSConnection.class);

        Map<String, Object> headers = new HashMap<String, Object>();
        headers.put(MessageContext.HTTP_REQUEST_HEADERS, new HashMap<String, List<String>>());
        headers.put(WSConnection.RESPONSE_CODE, responseCode);

        Dispatch<DataSource> dispatch = Mockito.mock(Dispatch.class);
        Mockito.when(dispatch.getRequestContext()).thenReturn(headers);
        Mockito.when(dispatch.getResponseContext()).thenReturn(headers);

        Mockito.when(connection.createDispatch(Mockito.eq(HTTPBinding.HTTP_BINDING), Mockito.anyString(),
                Mockito.eq(DataSource.class), Mockito.eq(Service.Mode.MESSAGE), Mockito.anyBoolean())).thenReturn(dispatch);

        DataSource ds = new DataSource() {
            @Override
            public OutputStream getOutputStream() throws IOException {
                return new ByteArrayOutputStream();
            }

            @Override
            public String getName() {
                return "result";
            }

            @Override
            public InputStream getInputStream() throws IOException {
                ByteArrayInputStream in = new ByteArrayInputStream(resultJson.getBytes());
                return in;
            }

            @Override
            public String getContentType() {
                return "application/xml";
            }
        };
        Mockito.when(dispatch.invoke(Mockito.any())).thenReturn(ds);

        ProcedureExecution execution = translator
                .createProcedureExecution((Call) cmd, context,
                        utility.createRuntimeMetadata(), connection);
        execution.execute();

        ArgumentCaptor<String> endpoint = ArgumentCaptor.forClass(String.class);
        ArgumentCaptor<String> binding = ArgumentCaptor.forClass(String.class);

        Mockito.verify(connection).createDispatch(binding.capture(),
                endpoint.capture(), Mockito.eq(DataSource.class),
                Mockito.eq(Service.Mode.MESSAGE), Mockito.anyBoolean());
        assertEquals(expectedURL, decode ? URLDecoder.decode(endpoint.getValue(), "utf-8") : endpoint.getValue());
        return execution;
    }

    protected UpdateExecution helpUpdateExecute(MetadataFactory mf, String query, String expectedPayload,
                                                final String resultJson, String expectedURL, String expectedMethod, int responseCode)
            throws Exception {

        ODataExecutionFactory translator = new ODataExecutionFactory();
        translator.start();

        TranslationUtility utility = new TranslationUtility(
                TestODataMetadataProcessor.getTransformationMetadata(mf, translator));

        Command cmd = utility.parseCommand(query);
        ExecutionContext context = Mockito.mock(ExecutionContext.class);
        WSConnection connection = Mockito.mock(WSConnection.class);

        Map<String, Object> headers = new HashMap<String, Object>();
        headers.put(MessageContext.HTTP_REQUEST_HEADERS, new HashMap<String, List<String>>());
        headers.put(WSConnection.RESPONSE_CODE, responseCode);

        Dispatch<DataSource> dispatch = Mockito.mock(Dispatch.class);
        Mockito.when(dispatch.getRequestContext()).thenReturn(headers);
        Mockito.when(dispatch.getResponseContext()).thenReturn(headers);

        Mockito.when(connection.createDispatch(Mockito.eq(HTTPBinding.HTTP_BINDING), Mockito.anyString(),
                Mockito.eq(DataSource.class), Mockito.eq(Service.Mode.MESSAGE), Mockito.anyBoolean())).thenReturn(dispatch);

        DataSource ds = new DataSource() {
            @Override
            public OutputStream getOutputStream() throws IOException {
                return new ByteArrayOutputStream();
            }

            @Override
            public String getName() {
                return "result";
            }

            @Override
            public InputStream getInputStream() throws IOException {
                ByteArrayInputStream in = new ByteArrayInputStream(resultJson.getBytes());
                return in;
            }

            @Override
            public String getContentType() {
                return "application/json";
            }
        };
        ArgumentCaptor<DataSource> data = ArgumentCaptor.forClass(DataSource.class);
        Mockito.when(dispatch.invoke(data.capture())).thenReturn(ds);

        UpdateExecution execution = translator
                .createUpdateExecution(cmd, context,
                        utility.createRuntimeMetadata(), connection);
        execution.execute();

        ArgumentCaptor<String> endpoint = ArgumentCaptor.forClass(String.class);
        ArgumentCaptor<String> binding = ArgumentCaptor.forClass(String.class);


        Mockito.verify(connection).createDispatch(binding.capture(),
                endpoint.capture(), Mockito.eq(DataSource.class),
                Mockito.eq(Service.Mode.MESSAGE), Mockito.anyBoolean());
        assertEquals(expectedURL, URLDecoder.decode(endpoint.getValue(), "utf-8"));
        String payload = new String(ObjectConverterUtil.convertToByteArray(
                ((InputStreamFactory.ClobInputStreamFactory) data.getValue()).getInputStream()));
        assertEquals(expectedPayload, payload);
        assertEquals(expectedMethod, dispatch.getRequestContext().get(MessageContext.HTTP_REQUEST_METHOD));
        return execution;
    }
}
