package com.dexorder.flink.sink;

import com.dexorder.flink.publisher.MarketWrapper;
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;
import org.zeromq.SocketType;
import org.zeromq.ZContext;
import org.zeromq.ZMQ;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

/**
 * Writes symbol metadata to Iceberg symbol_metadata table.
 *
 * Deduplicates symbols in-memory to prevent writing duplicates.
 * Batches writes by exchange to minimize file fragmentation.
 * After committing to Iceberg, sends a notification via ZMQ PUSH socket to job manager.
 */
public class SymbolMetadataWriter extends RichFlatMapFunction<MarketWrapper, MarketWrapper> {
    private static final long serialVersionUID = 1L;
    private static final Logger LOG = LoggerFactory.getLogger(SymbolMetadataWriter.class);
    private static final int BATCH_SIZE = 500; // Commit every 500 symbols per exchange
    private static final byte PROTOCOL_VERSION = 0x01;
    private static final byte MSG_TYPE_SYMBOL_METADATA_UPDATED = 0x13;

    private final TableLoader tableLoader;
    private final String notificationEndpoint; // Job manager's PULL socket endpoint

    private transient Table table;
    private transient Set<String> seenSymbols; // Track seen symbols to prevent duplicates
    private transient GenericAppenderFactory appenderFactory;
    private transient OutputFileFactory fileFactory;
    private transient ZContext zmqContext;
    private transient ZMQ.Socket pushSocket; // PUSH socket to job manager

    // Batching state per exchange
    private transient java.util.Map<String, DataWriter<Record>> writersByExchange;
    private transient java.util.Map<String, Integer> countsPerExchange;
    private transient java.util.Map<String, List<MarketWrapper>> pendingOutputPerExchange;

    public SymbolMetadataWriter(TableLoader tableLoader, String notificationEndpoint) {
        this.tableLoader = tableLoader;
        this.notificationEndpoint = notificationEndpoint;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        tableLoader.open();
        table = tableLoader.loadTable();
        seenSymbols = new HashSet<>();
        writersByExchange = new java.util.HashMap<>();
        countsPerExchange = new java.util.HashMap<>();
        pendingOutputPerExchange = new java.util.HashMap<>();

        // Initialize ZMQ PUSH socket to job manager (mimics HistoryNotificationPublisher pattern)
        zmqContext = new ZContext();
        pushSocket = zmqContext.createSocket(SocketType.PUSH);
        pushSocket.setLinger(1000);
        pushSocket.setSndHWM(10000);
        pushSocket.connect(notificationEndpoint);

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

        LOG.info("SymbolMetadataWriter opened, table loaded: {}", table.name());
        LOG.info("Connected PUSH socket to notification endpoint: {}", notificationEndpoint);
    }

    @Override
    public void flatMap(MarketWrapper market, Collector<MarketWrapper> out) throws Exception {
        // Create unique key for deduplication (internal key, not stored)
        String symbolKey = market.getMarketId() + "." + market.getExchangeId();

        // Skip if we've already seen this symbol
        if (seenSymbols.contains(symbolKey)) {
            LOG.debug("Skipping duplicate symbol: {}", symbolKey);
            out.collect(market); // Still emit for downstream processing
            return;
        }

        // Mark as seen
        seenSymbols.add(symbolKey);

        String exchangeId = market.getExchangeId();

        // Create Iceberg record from Market protobuf
        GenericRecord record = GenericRecord.create(table.schema());
        record.setField("exchange_id", exchangeId);
        record.setField("market_id", market.getMarketId());
        record.setField("market_type", market.getMarketType());
        record.setField("description", market.getDescription());
        record.setField("base_asset", market.getBaseAsset());
        record.setField("quote_asset", market.getQuoteAsset());

        // Convert supported_period_seconds to List<Integer>
        List<Integer> supportedPeriods = new ArrayList<>(market.getSupportedPeriodSeconds());
        record.setField("supported_period_seconds", supportedPeriods);

        record.setField("earliest_time", market.getEarliestTime() != 0 ? market.getEarliestTime() : null);
        record.setField("updated_at", System.currentTimeMillis() * 1_000_000L); // Current time in nanoseconds

        // Nautilus Instrument fields (populated from exchange API data)
        record.setField("price_precision", market.getPricePrecision() != 0 ? market.getPricePrecision() : null);
        record.setField("size_precision", market.getSizePrecision() != 0 ? market.getSizePrecision() : null);
        record.setField("tick_size", market.getTickSize() != 0.0 ? market.getTickSize() : null);
        record.setField("lot_size", market.getLotSize() != 0.0 ? market.getLotSize() : null);
        record.setField("min_notional", market.getMinNotional() != 0.0 ? market.getMinNotional() : null);
        record.setField("margin_init", market.getMarginInit() != 0.0 ? market.getMarginInit() : null);
        record.setField("margin_maint", market.getMarginMaint() != 0.0 ? market.getMarginMaint() : null);
        record.setField("maker_fee", market.getMakerFee() != 0.0 ? market.getMakerFee() : null);
        record.setField("taker_fee", market.getTakerFee() != 0.0 ? market.getTakerFee() : null);
        record.setField("contract_multiplier", market.getContractMultiplier() != 0.0 ? market.getContractMultiplier() : null);

        // Get or create writer for this exchange
        DataWriter<Record> writer = writersByExchange.get(exchangeId);
        if (writer == null) {
            // Compute partition key from exchange_id
            GenericRecord partitionRecord = GenericRecord.create(table.schema());
            partitionRecord.setField("exchange_id", exchangeId);
            PartitionKey partitionKey = new PartitionKey(table.spec(), table.schema());
            partitionKey.partition(partitionRecord);

            // Create new writer for this exchange's partition
            EncryptedOutputFile encryptedFile = fileFactory.newOutputFile(partitionKey);
            writer = appenderFactory.newDataWriter(encryptedFile, FileFormat.PARQUET, partitionKey);
            writersByExchange.put(exchangeId, writer);
            countsPerExchange.put(exchangeId, 0);
            pendingOutputPerExchange.put(exchangeId, new ArrayList<>());
        }

        // Write record to batch
        writer.write(record);

        // Track count and pending output
        int count = countsPerExchange.get(exchangeId) + 1;
        countsPerExchange.put(exchangeId, count);
        pendingOutputPerExchange.get(exchangeId).add(market);

        // Flush batch if we've reached the batch size
        if (count >= BATCH_SIZE) {
            flushExchange(exchangeId, out);
        }
    }

    private void flushExchange(String exchangeId, Collector<MarketWrapper> out) throws Exception {
        DataWriter<Record> writer = writersByExchange.get(exchangeId);
        if (writer == null) {
            return;
        }

        try {
            writer.close();
            table.newAppend()
                    .appendFile(writer.toDataFile())
                    .commit();

            int count = countsPerExchange.get(exchangeId);
            LOG.info("Committed batch of {} symbols for exchange: {}", count, exchangeId);

            // Send notification to gateways to reload symbol metadata
            sendMetadataUpdateNotification(exchangeId, count);

            // Emit all pending outputs
            for (MarketWrapper market : pendingOutputPerExchange.get(exchangeId)) {
                out.collect(market);
            }
        } finally {
            // Clear state for this exchange
            writersByExchange.remove(exchangeId);
            countsPerExchange.remove(exchangeId);
            pendingOutputPerExchange.remove(exchangeId);
        }
    }

    private void sendMetadataUpdateNotification(String exchangeId, int count) {
        if (pushSocket == null) {
            LOG.warn("Push socket is null - cannot send METADATA_UPDATE notification for exchange: {} ({} symbols)", exchangeId, count);
            return;
        }

        try {
            // Topic for metadata updates (broadcast to all gateways)
            String topic = "METADATA_UPDATE";

            // Empty payload (notification only, no additional data needed)
            byte[] payload = new byte[0];

            // Build message frame: [msg_type][payload]
            byte[] messageFrame = new byte[1 + payload.length];
            messageFrame[0] = MSG_TYPE_SYMBOL_METADATA_UPDATED;
            System.arraycopy(payload, 0, messageFrame, 1, payload.length);

            // Send three-frame message via PUSH: [topic][version][message]
            // Job manager's forwarder will republish via MARKET_DATA_PUB
            pushSocket.sendMore(topic);
            pushSocket.sendMore(new byte[]{PROTOCOL_VERSION});
            pushSocket.send(messageFrame, 0);

            LOG.info("Sent METADATA_UPDATE notification for exchange: {} ({} symbols)", exchangeId, count);
        } catch (Exception e) {
            LOG.error("Failed to send metadata update notification for exchange: {}", exchangeId, e);
            // Don't throw - notification is best-effort
        }
    }

    @Override
    public void close() throws Exception {
        // Flush any remaining batches
        try {
            for (String exchangeId : new ArrayList<>(writersByExchange.keySet())) {
                DataWriter<Record> writer = writersByExchange.get(exchangeId);
                if (writer != null) {
                    try {
                        writer.close();
                        table.newAppend()
                                .appendFile(writer.toDataFile())
                                .commit();

                        int count = countsPerExchange.get(exchangeId);
                        LOG.info("Final flush: committed {} remaining symbols for exchange: {}", count, exchangeId);

                        // Send final notification
                        sendMetadataUpdateNotification(exchangeId, count);
                    } catch (Exception e) {
                        LOG.error("Failed to flush remaining batch for exchange: {}", exchangeId, e);
                    }
                }
            }
        } finally {
            writersByExchange.clear();
            countsPerExchange.clear();
            pendingOutputPerExchange.clear();
        }

        // Close ZMQ resources
        if (pushSocket != null) {
            pushSocket.close();
        }
        if (zmqContext != null) {
            zmqContext.close();
        }

        if (tableLoader != null) {
            tableLoader.close();
        }
    }
}
