package com.dexorder.flink.sink;

import com.dexorder.flink.publisher.OHLCBatchWrapper;
import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.util.Collector;
import org.apache.iceberg.FileFormat;
import org.apache.iceberg.PartitionKey;
import org.apache.iceberg.Table;
import org.apache.iceberg.data.GenericAppenderFactory;
import org.apache.iceberg.data.GenericRecord;
import org.apache.iceberg.data.Record;
import org.apache.iceberg.encryption.EncryptedOutputFile;
import org.apache.iceberg.flink.TableLoader;
import org.apache.iceberg.io.DataWriter;
import org.apache.iceberg.io.OutputFileFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Writes historical OHLC batches directly to Iceberg using the catalog API.
 *
 * Unlike the streaming sink (FlinkSink), this uses table.newAppend().commit() which
 * commits synchronously and immediately — no checkpoint dependency. Batches are emitted
 * downstream only after the commit returns, guaranteeing HistoryNotificationFunction
 * fires after data is visible to readers.
 */
public class HistoricalBatchWriter extends RichFlatMapFunction<OHLCBatchWrapper, OHLCBatchWrapper> {
    private static final long serialVersionUID = 1L;
    private static final Logger LOG = LoggerFactory.getLogger(HistoricalBatchWriter.class);

    private final TableLoader tableLoader;
    private transient Table table;

    public HistoricalBatchWriter(TableLoader tableLoader) {
        this.tableLoader = tableLoader;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        tableLoader.open();
        table = tableLoader.loadTable();
        LOG.info("HistoricalBatchWriter opened, table loaded: {}", table.name());
    }

    @Override
    public void flatMap(OHLCBatchWrapper batch, Collector<OHLCBatchWrapper> out) throws Exception {
        // Empty batches (NOT_FOUND/ERROR markers): emit immediately without writing
        if (batch.getRows() == null || batch.getRows().isEmpty()) {
            LOG.debug("Passing through empty batch (marker): request_id={}, status={}",
                    batch.getRequestId(), batch.getStatus());
            out.collect(batch);
            return;
        }

        GenericAppenderFactory appenderFactory = new GenericAppenderFactory(table.schema(), table.spec());
        OutputFileFactory fileFactory = OutputFileFactory
                .builderFor(table, getRuntimeContext().getIndexOfThisSubtask(), System.nanoTime())
                .format(FileFormat.PARQUET)
                .build();

        // Compute partition key from ticker (all rows in a batch share one ticker)
        GenericRecord partitionRecord = GenericRecord.create(table.schema());
        partitionRecord.setField("ticker", batch.getTicker());
        PartitionKey partitionKey = new PartitionKey(table.spec(), table.schema());
        partitionKey.partition(partitionRecord);

        // Write all rows to one data file
        EncryptedOutputFile encryptedFile = fileFactory.newOutputFile(partitionKey);
        DataWriter<Record> writer = appenderFactory.newDataWriter(
                encryptedFile, FileFormat.PARQUET, partitionKey);

        long ingestedAt = System.currentTimeMillis() * 1000;
        try {
            for (OHLCBatchWrapper.OHLCRow row : batch.getRows()) {
                GenericRecord record = GenericRecord.create(table.schema());
                record.setField("ticker", batch.getTicker());
                record.setField("period_seconds", batch.getPeriodSeconds());
                record.setField("timestamp", row.getTimestamp());
                record.setField("open", row.getOpen());
                record.setField("high", row.getHigh());
                record.setField("low", row.getLow());
                record.setField("close", row.getClose());
                record.setField("volume", row.getVolume() != 0 ? row.getVolume() : null);
                record.setField("buy_vol", null);
                record.setField("sell_vol", null);
                record.setField("open_time", null);
                record.setField("high_time", null);
                record.setField("low_time", null);
                record.setField("close_time", null);
                record.setField("open_interest", null);
                record.setField("request_id", batch.getRequestId());
                record.setField("ingested_at", ingestedAt);
                writer.write(record);
            }
        } finally {
            writer.close();
        }

        // Immediate commit — no checkpoint needed
        table.newAppend()
                .appendFile(writer.toDataFile())
                .commit();

        LOG.info("Committed {} rows to Iceberg for request_id={}", batch.getRowCount(), batch.getRequestId());

        // Emit batch downstream only after successful commit
        out.collect(batch);
    }

    @Override
    public void close() throws Exception {
        if (tableLoader != null) {
            tableLoader.close();
        }
    }
}
