package com.dexorder.flink.publisher;

import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.MapState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
import org.apache.flink.api.common.typeinfo.PrimitiveArrayTypeInfo;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.util.Collector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Accumulates ticks into OHLC bars for each configured period.
 *
 * Keyed by ticker. Maintains per-period accumulators in MapState.
 * Uses a "lazy boundary" approach: a new window is detected when a tick arrives after
 * the previous window's end time (based on processing clock). The completed bar is
 * emitted immediately when the boundary is crossed, so bars are delayed by at most
 * one tick interval (~10s for realtime polling).
 *
 * Periods are configurable at construction time. All configured periods are computed
 * for every ticker receiving ticks; the ZMQ publisher filters to active subscriptions.
 *
 * Accumulator layout (long[7]):
 *   [0] open
 *   [1] high
 *   [2] low
 *   [3] close
 *   [4] volume (sum of base amount)
 *   [5] windowStartMs (epoch ms)
 *   [6] tickCount
 */
public class RealtimeBarFunction extends RichFlatMapFunction<TickWrapper, RealtimeBar> {
    private static final Logger LOG = LoggerFactory.getLogger(RealtimeBarFunction.class);
    private static final long serialVersionUID = 1L;

    private final int[] periods;
    private transient MapState<Integer, long[]> accumState;

    /**
     * @param periods  Period lengths in seconds (e.g., 60, 300, 900, 3600)
     */
    public RealtimeBarFunction(int[] periods) {
        this.periods = periods;
    }

    @Override
    public void open(Configuration parameters) {
        MapStateDescriptor<Integer, long[]> desc = new MapStateDescriptor<>(
                "ohlcAccum",
                BasicTypeInfo.INT_TYPE_INFO,
                PrimitiveArrayTypeInfo.LONG_PRIMITIVE_ARRAY_TYPE_INFO
        );
        accumState = getRuntimeContext().getMapState(desc);
    }

    @Override
    public void flatMap(TickWrapper tick, Collector<RealtimeBar> out) throws Exception {
        if (tick == null) return;

        long nowMs = System.currentTimeMillis();

        for (int period : periods) {
            long periodMs = period * 1000L;
            long windowStart = (nowMs / periodMs) * periodMs;

            long[] accum = accumState.get(period);

            if (accum == null) {
                // First tick for this period
                accumState.put(period, openWindow(tick, windowStart));

            } else if (accum[5] != windowStart) {
                // Window boundary crossed — emit completed bar then start fresh
                if (accum[6] > 0) {
                    out.collect(toBar(tick.getTicker(), period, accum));
                    LOG.debug("Emitted bar: ticker={}, period={}s, windowStart={}, ticks={}",
                            tick.getTicker(), period, accum[5], accum[6]);
                }
                accumState.put(period, openWindow(tick, windowStart));

            } else {
                // Same window — update
                accum[1] = Math.max(accum[1], tick.getPrice()); // high
                accum[2] = Math.min(accum[2], tick.getPrice()); // low
                accum[3] = tick.getPrice();                      // close
                accum[4] += tick.getAmount();                    // volume
                accum[6]++;                                       // tick count
                accumState.put(period, accum);
            }
        }
    }

    private static long[] openWindow(TickWrapper tick, long windowStart) {
        return new long[]{
            tick.getPrice(), // open
            tick.getPrice(), // high
            tick.getPrice(), // low
            tick.getPrice(), // close
            tick.getAmount(), // volume
            windowStart,
            1L               // tickCount
        };
    }

    private static RealtimeBar toBar(String ticker, int period, long[] accum) {
        return new RealtimeBar(
                ticker, period,
                accum[5],           // windowStartMs
                accum[0], accum[1], accum[2], accum[3], // O H L C
                accum[4],           // volume
                (int) accum[6]      // tickCount
        );
    }
}
