/*
 * 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.fluss.lake.paimon.tiering;

import org.apache.fluss.lake.paimon.tiering.append.AppendOnlyWriter;
import org.apache.fluss.lake.paimon.tiering.mergetree.MergeTreeWriter;
import org.apache.fluss.lake.writer.LakeWriter;
import org.apache.fluss.lake.writer.WriterInitContext;
import org.apache.fluss.metadata.TablePath;
import org.apache.fluss.record.LogRecord;

import org.apache.paimon.CoreOptions;
import org.apache.paimon.catalog.Catalog;
import org.apache.paimon.table.FileStoreTable;
import org.apache.paimon.table.sink.CommitMessage;

import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.Map;

import static org.apache.fluss.lake.paimon.utils.PaimonConversions.toPaimon;

/** Implementation of {@link LakeWriter} for Paimon. */
public class PaimonLakeWriter implements LakeWriter<PaimonWriteResult> {

    private final Catalog paimonCatalog;
    private final RecordWriter<?> recordWriter;

    public PaimonLakeWriter(
            PaimonCatalogProvider paimonCatalogProvider, WriterInitContext writerInitContext)
            throws IOException {
        this.paimonCatalog = paimonCatalogProvider.get();
        FileStoreTable fileStoreTable =
                getTable(
                        writerInitContext.tablePath(),
                        writerInitContext.tableInfo().getTableConfig().isDataLakeAutoCompaction());

        List<String> partitionKeys = fileStoreTable.partitionKeys();

        this.recordWriter =
                fileStoreTable.primaryKeys().isEmpty()
                        ? new AppendOnlyWriter(
                                fileStoreTable,
                                writerInitContext.tableBucket(),
                                writerInitContext.partition(),
                                partitionKeys)
                        : new MergeTreeWriter(
                                fileStoreTable,
                                writerInitContext.tableBucket(),
                                writerInitContext.partition(),
                                partitionKeys);
    }

    @Override
    public void write(LogRecord record) throws IOException {
        try {
            recordWriter.write(record);
        } catch (Exception e) {
            throw new IOException("Failed to write Fluss record to Paimon.", e);
        }
    }

    @Override
    public PaimonWriteResult complete() throws IOException {
        CommitMessage commitMessage;
        try {
            commitMessage = recordWriter.complete();
        } catch (Exception e) {
            throw new IOException("Failed to complete Paimon write.", e);
        }
        return new PaimonWriteResult(commitMessage);
    }

    @Override
    public void close() throws IOException {
        try {
            if (recordWriter != null) {
                recordWriter.close();
            }
            if (paimonCatalog != null) {
                paimonCatalog.close();
            }
        } catch (Exception e) {
            throw new IOException("Failed to close PaimonLakeWriter.", e);
        }
    }

    private FileStoreTable getTable(TablePath tablePath, boolean isAutoCompaction)
            throws IOException {
        try {
            FileStoreTable table = (FileStoreTable) paimonCatalog.getTable(toPaimon(tablePath));
            Map<String, String> compactionOptions =
                    Collections.singletonMap(
                            CoreOptions.WRITE_ONLY.key(),
                            isAutoCompaction ? Boolean.FALSE.toString() : Boolean.TRUE.toString());
            return table.copy(compactionOptions);
        } catch (Exception e) {
            throw new IOException("Failed to get table " + tablePath + " in Paimon.", e);
        }
    }
}
