package com.dexorder.flink;

import com.dexorder.flink.config.AppConfig;
import com.dexorder.flink.iceberg.SchemaInitializer;
import com.dexorder.flink.ingestor.IngestorControlChannel;
import com.dexorder.flink.ingestor.IngestorWorkQueue;
import com.dexorder.flink.kafka.TopicManager;
import com.dexorder.flink.publisher.HistoryNotificationForwarder;
import com.dexorder.flink.publisher.HistoryNotificationFunction;
import com.dexorder.flink.publisher.OHLCBatchWrapper;
import com.dexorder.flink.publisher.OHLCBatchDeserializer;
import com.dexorder.flink.sink.HistoricalBatchWriter;
import com.dexorder.flink.zmq.ZmqChannelManager;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.streaming.api.functions.sink.DiscardingSink;
import org.apache.flink.connector.kafka.source.KafkaSource;
import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.hadoop.conf.Configuration;
import org.apache.iceberg.catalog.TableIdentifier;
import org.apache.iceberg.flink.CatalogLoader;
import org.apache.iceberg.flink.TableLoader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

/**
 * Main entry point for the Trading Flink application.
 *
 * Responsibilities:
 * - Load configuration and secrets
 * - Initialize ZMQ channels for ingestor communication, market data pub/sub, and client requests
 * - Set up Kafka connectors for data ingestion
 * - Bootstrap the Flink streaming job
 */
public class TradingFlinkApp {
    private static final Logger LOG = LoggerFactory.getLogger(TradingFlinkApp.class);

    private static final String DEFAULT_CONFIG_PATH = "/etc/config/config.yaml";
    private static final String DEFAULT_SECRETS_PATH = "/etc/secrets/secrets.yaml";

    public static void main(String[] args) throws Exception {
        LOG.info("Starting Trading Flink Application");

        // Load configuration
        String configPath = System.getenv().getOrDefault("CONFIG_PATH", DEFAULT_CONFIG_PATH);
        String secretsPath = System.getenv().getOrDefault("SECRETS_PATH", DEFAULT_SECRETS_PATH);

        AppConfig config = loadConfig(configPath, secretsPath);
        LOG.info("Configuration loaded successfully");

        // Initialize Kafka topics
        try (TopicManager topicManager = new TopicManager(config.getKafkaBootstrapServers())) {
            topicManager.ensureTopicsExist();
            LOG.info("Kafka topics initialized: {}", topicManager.getTopicNames());
        } catch (Exception e) {
            LOG.error("Failed to initialize Kafka topics", e);
            throw e;
        }

        // Initialize Iceberg schemas
        try {
            Map<String, String> catalogProps = new HashMap<>();
            catalogProps.put("type", "rest");
            catalogProps.put("uri", config.getString("iceberg_catalog_uri", "http://iceberg-catalog:8181"));
            catalogProps.put("warehouse", config.getString("iceberg_warehouse", "s3://warehouse/"));

            // Configure S3 for MinIO
            catalogProps.put("s3.endpoint", config.getString("s3_endpoint", "http://minio:9000"));
            catalogProps.put("s3.path-style-access", "true");
            catalogProps.put("client.region", "us-east-1");

            // Use CatalogLoader.rest() for REST catalog instead of custom loader
            CatalogLoader catalogLoader = CatalogLoader.rest(
                    "trading",
                    new Configuration(),
                    catalogProps
            );

            org.apache.iceberg.catalog.Catalog catalog = catalogLoader.loadCatalog();
            try {
                SchemaInitializer schemaInitializer = new SchemaInitializer(
                        catalog,
                        config.getIcebergNamespace()
                );
                schemaInitializer.initializeSchemas();
            } finally {
                if (catalog instanceof java.io.Closeable) {
                    ((java.io.Closeable) catalog).close();
                }
            }
            LOG.info("Iceberg schemas initialized");
        } catch (Exception e) {
            LOG.error("Failed to initialize Iceberg schemas", e);
            throw e;
        }

        // Initialize ZeroMQ channels
        ZmqChannelManager zmqManager = new ZmqChannelManager(config);
        try {
            zmqManager.initializeChannels();
            LOG.info("ZeroMQ channels initialized");

            // Initialize history notification forwarder (runs in job manager)
            // Binds PULL socket to receive notifications from task managers, forwards to MARKET_DATA_PUB
            HistoryNotificationForwarder notificationForwarder = new HistoryNotificationForwarder(
                    config.getNotificationPullPort(),
                    zmqManager.getSocket(ZmqChannelManager.Channel.MARKET_DATA_PUB)
            );
            notificationForwarder.start();
            LOG.info("History notification forwarder started on port {}", config.getNotificationPullPort());

            // Initialize ingestor components
            IngestorWorkQueue workQueue = new IngestorWorkQueue(zmqManager);
            IngestorControlChannel controlChannel = new IngestorControlChannel(zmqManager);

            // Start the work queue processor
            workQueue.start();
            LOG.info("Ingestor work queue started");

            // Set up Flink streaming environment
            StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

            // Configure Flink
            configureFlinkEnvironment(env, config);

            // Set up Kafka source for OHLCBatch data
            KafkaSource<OHLCBatchWrapper> ohlcSource = KafkaSource.<OHLCBatchWrapper>builder()
                    .setBootstrapServers(config.getKafkaBootstrapServers())
                    .setTopics(config.getKafkaOhlcTopic())
                    .setGroupId("flink-ohlc-consumer")
                    .setStartingOffsets(OffsetsInitializer.earliest())
                    .setValueOnlyDeserializer(new OHLCBatchDeserializer())
                    .build();

            // Create OHLCBatch data stream
            DataStream<OHLCBatchWrapper> ohlcStream = env
                    .fromSource(ohlcSource, WatermarkStrategy.noWatermarks(), "OHLCBatch Kafka Source");

            LOG.info("OHLCBatch Kafka source configured");

            // Set up Iceberg catalog and table loader
            Map<String, String> catalogProps2 = new HashMap<>();
            catalogProps2.put("type", "rest");
            catalogProps2.put("uri", config.getString("iceberg_catalog_uri", "http://iceberg-catalog:8181"));
            catalogProps2.put("warehouse", config.getString("iceberg_warehouse", "s3://warehouse/"));

            // Configure S3 for MinIO
            catalogProps2.put("s3.endpoint", config.getString("s3_endpoint", "http://minio:9000"));
            catalogProps2.put("s3.path-style-access", "true");
            catalogProps2.put("client.region", "us-east-1");

            CatalogLoader catalogLoader2 = CatalogLoader.rest(
                    "trading",
                    new Configuration(),
                    catalogProps2
            );

            TableLoader tableLoader = TableLoader.fromCatalog(
                    catalogLoader2,
                    TableIdentifier.of(config.getIcebergNamespace(), "ohlc")
            );

            LOG.info("Iceberg table loader configured: {}.ohlc", config.getIcebergNamespace());

            // Historical pipeline: write to Iceberg first, then notify.
            // HistoricalBatchWriter uses direct catalog API (table.newAppend().commit()),
            // which commits synchronously — no checkpoint dependency. Batches are emitted
            // downstream only after commit returns, so notifications are guaranteed to fire
            // after data is visible to readers.
            // Parallelism MUST be 1: HistoryNotificationPublisher binds a ZMQ PUB socket,
            // and only one instance can bind the same port.
            DataStream<OHLCBatchWrapper> processedStream = ohlcStream
                    .flatMap(new HistoricalBatchWriter(tableLoader))
                    .setParallelism(1)
                    .process(new HistoryNotificationFunction(
                            config.getNotificationPublishEndpoint(),
                            config.getIcebergNamespace(),
                            config.getIcebergTablePrefix()
                    ))
                    .setParallelism(1);

            // Add a discard sink to force Flink to include the pipeline in the execution graph.
            processedStream.addSink(new DiscardingSink<>()).setParallelism(1);

            LOG.info("Historical pipeline configured: HistoricalBatchWriter -> HistoryNotificationFunction");

            // TODO: Set up CEP patterns and triggers
            // TODO: Set up realtime tick processing

            LOG.info("Flink job configured, starting execution");

            // Register shutdown hook for cleanup
            Runtime.getRuntime().addShutdownHook(new Thread(() -> {
                LOG.info("Shutting down Trading Flink Application");
                try {
                    // Send shutdown signal to ingestors
                    controlChannel.shutdown();

                    // Stop work queue
                    workQueue.stop();

                    // Stop notification forwarder
                    notificationForwarder.close();

                    // Close ZMQ channels
                    zmqManager.close();

                    LOG.info("Shutdown complete");
                } catch (Exception e) {
                    LOG.error("Error during shutdown", e);
                }
            }));

            // Execute the Flink job
            env.execute("Trading Flink Application");

        } catch (Exception e) {
            LOG.error("Fatal error in Trading Flink Application", e);
            zmqManager.close();
            throw e;
        }
    }

    private static AppConfig loadConfig(String configPath, String secretsPath) throws IOException {
        LOG.info("Loading config from: {}", configPath);
        LOG.info("Loading secrets from: {}", secretsPath);

        try {
            return new AppConfig(configPath, secretsPath);
        } catch (IOException e) {
            LOG.error("Failed to load configuration", e);
            throw e;
        }
    }

    private static void configureFlinkEnvironment(StreamExecutionEnvironment env, AppConfig config) {
        // Set parallelism - defaults to 1; increase when more task slots are available
        env.setParallelism(config.getInt("flink_parallelism", 1));

        // Enable checkpointing for fault tolerance
        env.enableCheckpointing(60000); // Checkpoint every 60 seconds

        // Set time characteristic
        // env.setStreamTimeCharacteristic(TimeCharacteristic.EventTime); // Deprecated in newer Flink versions

        LOG.info("Flink environment configured");
    }
}
