/*
 * 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.iceberg.rest.responses;

import static org.apache.iceberg.TestBase.FILE_A;
import static org.apache.iceberg.TestBase.FILE_A_DELETES;
import static org.apache.iceberg.TestBase.PARTITION_SPECS_BY_ID;
import static org.apache.iceberg.TestBase.SCHEMA;
import static org.apache.iceberg.TestBase.SPEC;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;

import com.fasterxml.jackson.core.JsonFactory;
import com.fasterxml.jackson.core.JsonFactoryBuilder;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.InjectableValues;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.ObjectReader;
import java.util.List;
import org.apache.iceberg.BaseFileScanTask;
import org.apache.iceberg.DeleteFile;
import org.apache.iceberg.FileScanTask;
import org.apache.iceberg.PartitionSpecParser;
import org.apache.iceberg.SchemaParser;
import org.apache.iceberg.expressions.Expressions;
import org.apache.iceberg.expressions.ResidualEvaluator;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
import org.apache.iceberg.rest.PlanStatus;
import org.apache.iceberg.rest.RESTSerializers;
import org.apache.iceberg.rest.credentials.Credential;
import org.apache.iceberg.rest.credentials.ImmutableCredential;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

public class TestFetchPlanningResultResponseParser {

  private static final JsonFactory FACTORY =
      new JsonFactoryBuilder()
          .configure(JsonFactory.Feature.INTERN_FIELD_NAMES, false)
          .configure(JsonFactory.Feature.FAIL_ON_SYMBOL_HASH_OVERFLOW, false)
          .build();
  private static final ObjectMapper MAPPER = new ObjectMapper(FACTORY);

  @BeforeEach
  public void before() {
    RESTSerializers.registerAll(MAPPER);
  }

  @Test
  public void nullAndEmptyCheck() {
    assertThatThrownBy(() -> FetchPlanningResultResponseParser.toJson(null))
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessage("Invalid fetchPlanningResult response: null");

    assertThatThrownBy(
            () ->
                FetchPlanningResultResponseParser.fromJson(
                    (String) null, PARTITION_SPECS_BY_ID, false))
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessage("Invalid fetchPlanningResult response: null or empty");
  }

  @Test
  public void roundTripSerdeWithEmptyObject() {
    assertThatThrownBy(
            () ->
                FetchPlanningResultResponseParser.toJson(
                    FetchPlanningResultResponse.builder().build()))
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessage("Invalid status: null");

    String emptyJson = "{ }";
    assertThatThrownBy(
            () ->
                FetchPlanningResultResponseParser.fromJson(emptyJson, PARTITION_SPECS_BY_ID, false))
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessage("Invalid fetchPlanningResult response: null or empty");
  }

  @Test
  public void roundTripSerdeWithInvalidPlanStatus() {
    String invalidStatusJson = "{\"status\": \"someStatus\"}";
    assertThatThrownBy(
            () ->
                FetchPlanningResultResponseParser.fromJson(
                    invalidStatusJson, PARTITION_SPECS_BY_ID, false))
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessage("Invalid status name: someStatus");
  }

  @Test
  public void roundTripSerdeWithValidSubmittedStatus() {
    PlanStatus planStatus = PlanStatus.fromName("submitted");
    FetchPlanningResultResponse response =
        FetchPlanningResultResponse.builder().withPlanStatus(planStatus).build();

    String expectedJson = "{\"status\":\"submitted\"}";
    String json = FetchPlanningResultResponseParser.toJson(response);
    assertThat(json).isEqualTo(expectedJson);

    FetchPlanningResultResponse fromResponse =
        FetchPlanningResultResponseParser.fromJson(json, PARTITION_SPECS_BY_ID, false);
    assertThat(FetchPlanningResultResponseParser.toJson(fromResponse)).isEqualTo(expectedJson);
  }

  @Test
  public void roundTripSerdeWithInvalidPlanStatusSubmittedWithTasksPresent() {
    PlanStatus planStatus = PlanStatus.fromName("submitted");
    assertThatThrownBy(
            () ->
                FetchPlanningResultResponse.builder()
                    .withPlanStatus(planStatus)
                    .withPlanTasks(List.of("task1", "task2"))
                    .build())
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessage("Invalid response: tasks can only be returned in a 'completed' status");

    String invalidJson = "{\"status\":\"submitted\"," + "\"plan-tasks\":[\"task1\",\"task2\"]}";

    assertThatThrownBy(
            () ->
                FetchPlanningResultResponseParser.fromJson(
                    invalidJson, PARTITION_SPECS_BY_ID, false))
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessage("Invalid response: tasks can only be returned in a 'completed' status");
  }

  @Test
  public void roundTripSerdeWithInvalidPlanStatusSubmittedWithDeleteFilesNoFileScanTasksPresent() {
    PlanStatus planStatus = PlanStatus.fromName("submitted");
    assertThatThrownBy(
            () -> {
              FetchPlanningResultResponse.builder()
                  .withPlanStatus(planStatus)
                  .withDeleteFiles(List.of(FILE_A_DELETES))
                  .build();
            })
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessage(
            "Invalid response: deleteFiles should only be returned with fileScanTasks that reference them");

    String invalidJson =
        "{\"status\":\"submitted\","
            + "\"delete-files\":[{\"spec-id\":0,\"content\":\"position-deletes\","
            + "\"file-path\":\"/path/to/data-a-deletes.parquet\",\"file-format\":\"parquet\","
            + "\"partition\":[0],\"file-size-in-bytes\":10,\"record-count\":1}]"
            + "}";

    assertThatThrownBy(
            () ->
                FetchPlanningResultResponseParser.fromJson(
                    invalidJson, PARTITION_SPECS_BY_ID, false))
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessage(
            "Invalid response: deleteFiles should only be returned with fileScanTasks that reference them");
  }

  @Test
  public void roundTripSerdeWithValidStatusAndFileScanTasks() throws JsonProcessingException {
    ResidualEvaluator residualEvaluator =
        ResidualEvaluator.of(SPEC, Expressions.equal("id", 1), true);
    FileScanTask fileScanTask =
        new BaseFileScanTask(
            FILE_A,
            new DeleteFile[] {FILE_A_DELETES},
            SchemaParser.toJson(SCHEMA),
            PartitionSpecParser.toJson(SPEC),
            residualEvaluator);

    PlanStatus planStatus = PlanStatus.fromName("completed");
    FetchPlanningResultResponse response =
        FetchPlanningResultResponse.builder()
            .withPlanStatus(planStatus)
            .withFileScanTasks(List.of(fileScanTask))
            // assume this has been set
            .withSpecsById(PARTITION_SPECS_BY_ID)
            .build();

    String expectedToJson =
        "{\"status\":\"completed\","
            + "\"delete-files\":[{\"spec-id\":0,\"content\":\"position-deletes\","
            + "\"file-path\":\"/path/to/data-a-deletes.parquet\",\"file-format\":\"parquet\","
            + "\"partition\":[0],\"file-size-in-bytes\":10,\"record-count\":1}],"
            + "\"file-scan-tasks\":["
            + "{\"data-file\":{\"spec-id\":0,\"content\":\"data\",\"file-path\":\"/path/to/data-a.parquet\","
            + "\"file-format\":\"parquet\",\"partition\":[0],"
            + "\"file-size-in-bytes\":10,\"record-count\":1,\"sort-order-id\":0},"
            + "\"delete-file-references\":[0],"
            + "\"residual-filter\":{\"type\":\"eq\",\"term\":\"id\",\"value\":1}}]"
            + "}";

    String json = FetchPlanningResultResponseParser.toJson(response, false);
    assertThat(json).isEqualTo(expectedToJson);

    // use RESTObjectMapper to read this
    InjectableValues.Std injectableValues = new InjectableValues.Std();
    injectableValues.addValue("specsById", PARTITION_SPECS_BY_ID);
    injectableValues.addValue("caseSensitive", false);
    ObjectReader objectReader =
        MAPPER.readerFor(FetchPlanningResultResponse.class).with(injectableValues);
    FetchPlanningResultResponse fromResponse = objectReader.readValue(json);
    // Need to make a new response with partitionSpec set
    FetchPlanningResultResponse copyResponse =
        FetchPlanningResultResponse.builder()
            .withPlanStatus(fromResponse.planStatus())
            .withPlanTasks(fromResponse.planTasks())
            .withFileScanTasks(fromResponse.fileScanTasks())
            .withSpecsById(PARTITION_SPECS_BY_ID)
            .build();

    assertThat(FetchPlanningResultResponseParser.toJson(copyResponse, false))
        .isEqualTo(expectedToJson);
  }

  @Test
  public void emptyOrInvalidCredentials() {
    assertThat(
            FetchPlanningResultResponseParser.fromJson(
                    "{\"status\": \"completed\",\"storage-credentials\": null}",
                    PARTITION_SPECS_BY_ID,
                    false)
                .credentials())
        .isEmpty();

    assertThat(
            FetchPlanningResultResponseParser.fromJson(
                    "{\"status\": \"completed\",\"storage-credentials\": []}",
                    PARTITION_SPECS_BY_ID,
                    false)
                .credentials())
        .isEmpty();

    assertThatThrownBy(
            () ->
                FetchPlanningResultResponseParser.fromJson(
                    "{\"status\": \"completed\",\"storage-credentials\": \"invalid\"}",
                    PARTITION_SPECS_BY_ID,
                    false))
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessage("Cannot parse credentials from non-array: \"invalid\"");
  }

  @Test
  public void roundTripSerdeWithCredentials() {
    List<Credential> credentials =
        ImmutableList.of(
            ImmutableCredential.builder()
                .prefix("s3://custom-uri")
                .config(
                    ImmutableMap.of(
                        "s3.access-key-id",
                        "keyId",
                        "s3.secret-access-key",
                        "accessKey",
                        "s3.session-token",
                        "sessionToken"))
                .build(),
            ImmutableCredential.builder()
                .prefix("gs://custom-uri")
                .config(
                    ImmutableMap.of(
                        "gcs.oauth2.token", "gcsToken1", "gcs.oauth2.token-expires-at", "1000"))
                .build(),
            ImmutableCredential.builder()
                .prefix("gs")
                .config(
                    ImmutableMap.of(
                        "gcs.oauth2.token", "gcsToken2", "gcs.oauth2.token-expires-at", "2000"))
                .build());

    FetchPlanningResultResponse response =
        FetchPlanningResultResponse.builder()
            .withPlanStatus(PlanStatus.COMPLETED)
            .withCredentials(credentials)
            .build();

    String expectedJson =
        "{\n"
            + "  \"status\" : \"completed\",\n"
            + "  \"storage-credentials\" : [ {\n"
            + "    \"prefix\" : \"s3://custom-uri\",\n"
            + "    \"config\" : {\n"
            + "      \"s3.access-key-id\" : \"keyId\",\n"
            + "      \"s3.secret-access-key\" : \"accessKey\",\n"
            + "      \"s3.session-token\" : \"sessionToken\"\n"
            + "    }\n"
            + "  }, {\n"
            + "    \"prefix\" : \"gs://custom-uri\",\n"
            + "    \"config\" : {\n"
            + "      \"gcs.oauth2.token\" : \"gcsToken1\",\n"
            + "      \"gcs.oauth2.token-expires-at\" : \"1000\"\n"
            + "    }\n"
            + "  }, {\n"
            + "    \"prefix\" : \"gs\",\n"
            + "    \"config\" : {\n"
            + "      \"gcs.oauth2.token\" : \"gcsToken2\",\n"
            + "      \"gcs.oauth2.token-expires-at\" : \"2000\"\n"
            + "    }\n"
            + "  } ]\n"
            + "}";

    String json = FetchPlanningResultResponseParser.toJson(response, true);
    assertThat(json).isEqualTo(expectedJson);

    FetchPlanningResultResponse fromResponse =
        FetchPlanningResultResponseParser.fromJson(json, PARTITION_SPECS_BY_ID, false);
    FetchPlanningResultResponse copyResponse =
        FetchPlanningResultResponse.builder()
            .withPlanStatus(fromResponse.planStatus())
            .withCredentials(credentials)
            .build();

    assertThat(FetchPlanningResultResponseParser.toJson(copyResponse, true))
        .isEqualTo(expectedJson);
  }

  @Test
  public void roundTripSerdeWithFailedStatusAndErrorResponse() {
    ErrorResponse errorResponse =
        ErrorResponse.builder()
            .withMessage("Scan planning failed: table too large to plan")
            .withType("IllegalStateException")
            .responseCode(500)
            .build();

    FetchPlanningResultResponse response =
        FetchPlanningResultResponse.builder()
            .withPlanStatus(PlanStatus.FAILED)
            .withErrorResponse(errorResponse)
            .build();

    String expectedJson =
        "{\"status\":\"failed\","
            + "\"error\":{\"message\":\"Scan planning failed: table too large to plan\","
            + "\"type\":\"IllegalStateException\",\"code\":500}}";
    String json = FetchPlanningResultResponseParser.toJson(response);
    assertThat(json).isEqualTo(expectedJson);

    FetchPlanningResultResponse fromResponse =
        FetchPlanningResultResponseParser.fromJson(json, PARTITION_SPECS_BY_ID, false);
    assertThat(fromResponse.planStatus()).isEqualTo(PlanStatus.FAILED);
    assertThat(fromResponse.errorResponse()).isNotNull();
    assertThat(fromResponse.errorResponse().message())
        .isEqualTo("Scan planning failed: table too large to plan");
    assertThat(fromResponse.errorResponse().type()).isEqualTo("IllegalStateException");
    assertThat(fromResponse.errorResponse().code()).isEqualTo(500);
  }

  @Test
  public void parseFailedStatusWithoutErrorObject() {
    // Spec requires an `error` object on failed responses, but parse leniently so
    // a non-compliant server still surfaces the failure to the client.
    String json = "{\"status\":\"failed\"}";
    FetchPlanningResultResponse response =
        FetchPlanningResultResponseParser.fromJson(json, PARTITION_SPECS_BY_ID, false);
    assertThat(response.planStatus()).isEqualTo(PlanStatus.FAILED);
    assertThat(response.errorResponse()).isNull();
  }

  @Test
  public void parseFailedStatusWithPrimitiveErrorField() {
    String json = "{\"status\":\"failed\",\"error\":\"oops\"}";
    FetchPlanningResultResponse response =
        FetchPlanningResultResponseParser.fromJson(json, PARTITION_SPECS_BY_ID, false);
    assertThat(response.planStatus()).isEqualTo(PlanStatus.FAILED);
    assertThat(response.errorResponse()).isNull();
  }

  @Test
  public void cannotBuildWithErrorResponseWhenStatusIsNotFailed() {
    ErrorResponse errorResponse =
        ErrorResponse.builder().withMessage("boom").withType("X").responseCode(500).build();
    assertThatThrownBy(
            () ->
                FetchPlanningResultResponse.builder()
                    .withPlanStatus(PlanStatus.COMPLETED)
                    .withErrorResponse(errorResponse)
                    .build())
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessage("Invalid response: error can only be returned in a 'failed' status");
  }
}
