import java.math.BigDecimal;
import java.util.concurrent.ThreadLocalRandom;

import org.apache.fluss.client.Connection;
import org.apache.fluss.client.ConnectionFactory;
import org.apache.fluss.client.lookup.Lookuper;
import org.apache.fluss.client.lookup.LookupResult;
import org.apache.fluss.client.table.Table;
import org.apache.fluss.client.table.writer.UpsertWriter;
import org.apache.fluss.config.Configuration;
import org.apache.fluss.metadata.TablePath;
import org.apache.fluss.row.BinaryString;
import org.apache.fluss.row.Decimal;
import org.apache.fluss.row.GenericRow;
import org.apache.fluss.row.InternalRow;

// Run after completing the Apache Fluss Flink Quickstart.
//
//   MAVEN_OPTS='--add-opens=java.base/java.nio=ALL-UNNAMED' \
//     mvn compile exec:java -Dexec.mainClass=PrimaryKeyTableExample
//
// Set FLUSS_BOOTSTRAP when the Coordinator is not at localhost:9123.
public final class PrimaryKeyTableExample {
    private static final int CUSTOMER_ID = 999999;

    private PrimaryKeyTableExample() {}

    public static void main(String[] args) throws Exception {
        BigDecimal beforeBalance = randomBalanceExcept(null);
        BigDecimal afterBalance = randomBalanceExcept(beforeBalance);

        Configuration config = new Configuration();
        config.setString("bootstrap.servers", bootstrapAddress());

        try (Connection connection = ConnectionFactory.createConnection(config);
                Table profiles = connection.getTable(TablePath.of("demo", "customer_profile"))) {
            UpsertWriter writer = profiles.newUpsert().createWriter();
            Lookuper lookuper = profiles.newLookup().createLookuper();

            printProfile("[1/4] Initial value to write:", CUSTOMER_ID, "java-client-before", beforeBalance);
            writer.upsert(profileRow("java-client-before", beforeBalance)).get();
            writer.flush();

            printProfile("[2/4] Read the value after the first write:", lookup(lookuper));

            System.out.printf(
                    "[3/4] Change fields for the same primary key:%n"
                            + "membership=java-client-after%naccount_balance=%s%n",
                    afterBalance.toPlainString());
            writer.upsert(profileRow("java-client-after", afterBalance)).get();
            writer.flush();

            printProfile("[4/4] Read the value after the update:", lookup(lookuper));
        }
    }

    private static GenericRow profileRow(String membership, BigDecimal balance) {
        return GenericRow.of(
                CUSTOMER_ID,
                BinaryString.fromString("Quickstart User"),
                BinaryString.fromString(membership),
                Decimal.fromBigDecimal(balance, 15, 2));
    }

    private static InternalRow lookup(Lookuper lookuper) throws Exception {
        LookupResult result = lookuper.lookup(GenericRow.of(CUSTOMER_ID)).get();
        InternalRow row = result.getSingletonRow();
        if (row == null) {
            throw new IllegalStateException("customer " + CUSTOMER_ID + " was not found");
        }
        return row;
    }

    private static BigDecimal randomBalanceExcept(BigDecimal previous) {
        BigDecimal amount;
        do {
            amount = BigDecimal.valueOf(ThreadLocalRandom.current().nextLong(1_000, 100_000), 2);
        } while (amount.equals(previous));
        return amount;
    }

    private static void printProfile(String step, int customerId, String membership, BigDecimal balance) {
        System.out.printf(
                """
                %s
                {
                  "customer_id": %d,
                  "name": "Quickstart User",
                  "membership": "%s",
                  "account_balance": "%s"
                }
                %n""",
                step, customerId, membership, balance.toPlainString());
    }

    private static void printProfile(String step, InternalRow row) {
        printProfile(
                step,
                row.getInt(0),
                row.getString(2).toString(),
                row.getDecimal(3, 15, 2).toBigDecimal());
    }

    private static String bootstrapAddress() {
        String address = System.getenv("FLUSS_BOOTSTRAP");
        return address == null || address.isBlank() ? "localhost:9123" : address;
    }
}