/*
 * 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.kyuubi.engine.dataagent.prompt;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

import java.time.LocalDate;
import org.junit.jupiter.api.Test;

public class SystemPromptBuilderTest {

  @Test
  public void testDefaultBuildContainsBaseAndDate() {
    String prompt = SystemPromptBuilder.create().build();
    assertTrue(prompt.contains("data analysis agent"));
    assertTrue(prompt.contains(LocalDate.now().toString()));
  }

  @Test
  public void testPlaceholdersRemovedByDefault() {
    String prompt = SystemPromptBuilder.create().build();
    assertFalse(prompt.contains("{{tool_descriptions}}"));
  }

  @Test
  public void testToolDescriptionsSubstituted() {
    String tools = "- `sql_query`: Execute a SQL statement against the database.";
    String prompt = SystemPromptBuilder.create().toolDescriptions(tools).build();
    assertTrue(prompt.contains("sql_query"));
    assertTrue(prompt.contains("Execute a SQL statement"));
    assertFalse(prompt.contains("{{tool_descriptions}}"));
  }

  @Test
  public void testDatasourceSQLite() {
    String prompt = SystemPromptBuilder.create().datasource("sqlite").build();
    assertTrue(prompt.contains("SQLite SQL compatibility"));
    assertTrue(prompt.contains("JULIANDAY"));
  }

  @Test
  public void testDatasourceSpark() {
    String prompt = SystemPromptBuilder.create().datasource("spark").build();
    assertTrue(prompt.contains("data analysis agent"));
    assertTrue(prompt.contains("Spark SQL"));
  }

  @Test
  public void testWithCustomBase() {
    String prompt = SystemPromptBuilder.create().base("You are a helper.").build();
    assertTrue(prompt.startsWith("You are a helper."));
    assertTrue(prompt.contains(LocalDate.now().toString()));
  }

  @Test
  public void testWithSection() {
    String prompt = SystemPromptBuilder.create().section("Only query public schema.").build();
    assertTrue(prompt.contains("Only query public schema."));
  }

  @Test
  public void testFullComposition() {
    String prompt =
        SystemPromptBuilder.create()
            .toolDescriptions("- `sql_query`: Execute SQL.")
            .datasource("spark")
            .section("Limit all queries to 1000 rows.")
            .build();
    assertTrue(prompt.contains("sql_query"));
    assertTrue(prompt.contains("Spark SQL"));
    assertTrue(prompt.contains("Limit all queries to 1000 rows."));
    assertTrue(prompt.contains(LocalDate.now().toString()));
  }

  @Test
  public void testDatasourceReplacesNotAppends() {
    String prompt = SystemPromptBuilder.create().datasource("sqlite").datasource("spark").build();
    assertTrue(prompt.contains("Spark SQL"), "Last datasource wins");
    assertFalse(prompt.contains("SQLite SQL compatibility"), "Previous datasource replaced");
  }

  @Test
  public void testUnknownDatasourceUsesGenericDialectSection() {
    String prompt = SystemPromptBuilder.create().datasource("postgresql").build();
    assertTrue(prompt.contains("Current SQL dialect: postgresql"));
    assertTrue(prompt.contains("postgresql SQL syntax"));
  }

  @Test
  public void testDatasourceClickhouseUsesGenericDialectSection() {
    String prompt = SystemPromptBuilder.create().datasource("clickhouse").build();
    assertTrue(prompt.contains("Current SQL dialect: clickhouse"));
  }

  @Test
  public void testDatasourceMySQL() {
    String prompt = SystemPromptBuilder.create().datasource("mysql").build();
    assertTrue(prompt.contains("MySQL"));
  }

  @Test
  public void testLoadResourceThrowsOnMissing() {
    assertThrows(
        IllegalArgumentException.class,
        () -> {
          SystemPromptBuilder.loadResource("nonexistent-resource");
        });
  }

  @Test
  public void testNullsIgnored() {
    String plain = SystemPromptBuilder.create().build();
    assertEquals(plain, SystemPromptBuilder.create().datasource(null).build());
    assertEquals(plain, SystemPromptBuilder.create().toolDescriptions(null).build());
    assertEquals(plain, SystemPromptBuilder.create().section(null).build());
    assertEquals(plain, SystemPromptBuilder.create().section("").build());
  }
}
