/*
 * SPDX-License-Identifier: MIT
 */
package org.ta4j.core.backtest;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.WeakHashMap;
import org.ta4j.core.Bar;
import org.ta4j.core.BarSeries;
import org.ta4j.core.ExecutionSide;
import org.ta4j.core.Position;
import org.ta4j.core.Trade;
import org.ta4j.core.Trade.TradeType;
import org.ta4j.core.TradeFill;
import org.ta4j.core.TradingRecord;
import org.ta4j.core.backtest.ExecutionModelSupport.ExecutionTarget;
import org.ta4j.core.num.Num;

/**
 * Stop-limit execution model with partial fill progression.
 *
 * <p>
 * Strategy signals place stop-limit orders. Pending orders are evaluated on
 * each bar through {@link #onBar(int, TradingRecord, BarSeries)} and can be
 * filled progressively based on available bar volume participation. Partially
 * filled orders are committed on expiry when the target trading record supports
 * partial-exit accounting. Generated fills include execution side and bar end
 * timestamps so backtest fills match live-fill metadata shape.
 * </p>
 *
 * @since 0.22.4
 */
public class StopLimitExecutionModel implements TradeExecutionModel {

    private final Num stopTriggerRatio;
    private final Num limitOffsetRatio;
    private final Num maxBarParticipationRate;
    private final int maxBarsToFill;
    private final PriceSource priceSource;

    private final Map<TradingRecord, PendingOrder> pendingOrders = new WeakHashMap<>();
    private final Map<TradingRecord, List<RejectedOrder>> rejectedOrders = new WeakHashMap<>();

    /**
     * Creates a stop-limit execution model using next-bar open as the reference
     * price.
     *
     * @param stopTriggerRatio    stop trigger ratio in [0,1)
     * @param limitOffsetRatio    limit offset ratio in [0,1)
     * @param maxBarParticipation max per-bar fill participation in (0,1]
     * @param maxBarsToFill       order time-to-live in bars (>= 1)
     * @since 0.22.4
     */
    public StopLimitExecutionModel(Num stopTriggerRatio, Num limitOffsetRatio, Num maxBarParticipation,
            int maxBarsToFill) {
        this(stopTriggerRatio, limitOffsetRatio, maxBarParticipation, maxBarsToFill, PriceSource.NEXT_OPEN);
    }

    /**
     * Creates a stop-limit execution model.
     *
     * @param stopTriggerRatio    stop trigger ratio in [0,1)
     * @param limitOffsetRatio    limit offset ratio in [0,1) (must be >= stop
     *                            ratio)
     * @param maxBarParticipation max per-bar fill participation in (0,1]
     * @param maxBarsToFill       order time-to-live in bars (>= 1)
     * @param priceSource         base signal price source
     * @since 0.22.4
     */
    public StopLimitExecutionModel(Num stopTriggerRatio, Num limitOffsetRatio, Num maxBarParticipation,
            int maxBarsToFill, PriceSource priceSource) {
        validateRatio(stopTriggerRatio, "stopTriggerRatio");
        validateRatio(limitOffsetRatio, "limitOffsetRatio");
        validateRatio(maxBarParticipation, "maxBarParticipation");
        Objects.requireNonNull(priceSource, "priceSource");
        Num one = stopTriggerRatio.getNumFactory().one();
        if (stopTriggerRatio.isGreaterThanOrEqual(one)) {
            throw new IllegalArgumentException("stopTriggerRatio must be < 1");
        }
        if (limitOffsetRatio.isGreaterThanOrEqual(one)) {
            throw new IllegalArgumentException("limitOffsetRatio must be < 1");
        }
        if (maxBarParticipation.isZero()) {
            throw new IllegalArgumentException("maxBarParticipation must be > 0");
        }
        if (maxBarParticipation.isGreaterThan(one)) {
            throw new IllegalArgumentException("maxBarParticipation must be <= 1");
        }
        if (limitOffsetRatio.isLessThan(stopTriggerRatio)) {
            throw new IllegalArgumentException("limitOffsetRatio must be >= stopTriggerRatio");
        }
        if (maxBarsToFill < 1) {
            throw new IllegalArgumentException("maxBarsToFill must be >= 1");
        }
        this.stopTriggerRatio = stopTriggerRatio;
        this.limitOffsetRatio = limitOffsetRatio;
        this.maxBarParticipationRate = maxBarParticipation;
        this.maxBarsToFill = maxBarsToFill;
        this.priceSource = priceSource;
    }

    /**
     * Returns rejected orders for a trading record.
     *
     * @param tradingRecord trading record
     * @return rejected orders
     * @since 0.22.4
     */
    public List<RejectedOrder> getRejectedOrders(TradingRecord tradingRecord) {
        List<RejectedOrder> rejected = rejectedOrders.get(tradingRecord);
        return rejected == null ? List.of() : List.copyOf(rejected);
    }

    /**
     * Returns the current pending order snapshot for the trading record.
     *
     * @param tradingRecord trading record
     * @return pending order snapshot
     * @since 0.22.4
     */
    public Optional<PendingOrderSnapshot> getPendingOrder(TradingRecord tradingRecord) {
        PendingOrder order = pendingOrders.get(tradingRecord);
        if (order == null) {
            return Optional.empty();
        }
        return Optional.of(order.snapshot());
    }

    @Override
    public void execute(int index, TradingRecord tradingRecord, BarSeries barSeries, Num amount) {
        Objects.requireNonNull(tradingRecord, "tradingRecord");
        Objects.requireNonNull(barSeries, "barSeries");
        expireIfStale(index, tradingRecord);
        if (amount == null || amount.isNaN() || amount.isZero() || amount.isNegative()) {
            Num requestedAmount = amountOrZero(amount, barSeries);
            addRejectedOrder(tradingRecord,
                    new RejectedOrder(index, index, ExecutionModelSupport.nextTradeType(tradingRecord), requestedAmount,
                            requestedAmount.getNumFactory().zero(), "Invalid requested amount"));
            return;
        }
        Num requestedAmount = resolveRequestedAmount(tradingRecord, amount);
        PendingOrder pendingOrder = pendingOrders.get(tradingRecord);
        if (pendingOrder != null) {
            addRejectedOrder(tradingRecord,
                    new RejectedOrder(index, index, pendingOrder.tradeType, requestedAmount,
                            requestedAmount.getNumFactory().zero(),
                            "Signal ignored while another stop-limit order is pending"));
            return;
        }
        ExecutionTarget referenceTarget = ExecutionModelSupport.resolveExecutionTarget(index, barSeries, priceSource);
        if (referenceTarget == null) {
            addRejectedOrder(tradingRecord,
                    new RejectedOrder(index, index, ExecutionModelSupport.nextTradeType(tradingRecord), requestedAmount,
                            requestedAmount.getNumFactory().zero(),
                            "Unable to resolve reference bar for stop-limit order"));
            return;
        }

        TradeType tradeType = ExecutionModelSupport.nextTradeType(tradingRecord);
        Num stopPrice = toStopPrice(referenceTarget.price(), tradeType);
        Num limitPrice = toLimitPrice(referenceTarget.price(), tradeType);
        int activationIndex = resolveActivationIndex(referenceTarget.index());
        if (activationIndex > barSeries.getEndIndex()) {
            addRejectedOrder(tradingRecord, new RejectedOrder(index, index, tradeType, requestedAmount,
                    requestedAmount.getNumFactory().zero(), "Unable to resolve activation bar for stop-limit order"));
            return;
        }
        pendingOrders.put(tradingRecord, new PendingOrder(index, activationIndex, tradeType, requestedAmount, stopPrice,
                limitPrice, activationIndex + maxBarsToFill - 1));
    }

    @Override
    public void onBar(int index, TradingRecord tradingRecord, BarSeries barSeries) {
        PendingOrder order = pendingOrders.get(tradingRecord);
        if (order == null || index < order.activationIndex) {
            return;
        }

        Bar bar = barSeries.getBar(index);
        if (!order.triggered) {
            order.triggered = triggerReached(order.tradeType, bar, order.stopPrice);
        }

        if (order.triggered && limitReachable(order.tradeType, bar, order.limitPrice)) {
            Num fillAmount = fillAmount(order.remainingAmount(), bar.getVolume());
            if (fillAmount.isPositive()) {
                order.recordFill(index, bar, order.limitPrice, fillAmount);
            }
        }

        if (order.isCompletelyFilled()) {
            tradingRecord.operate(order.toTrade(tradingRecord));
            pendingOrders.remove(tradingRecord);
            return;
        }

        if (index >= order.expiryIndex) {
            expireOrder(index, tradingRecord, order);
        }
    }

    @Override
    public void onRunEnd(int lastProcessedIndex, TradingRecord tradingRecord) {
        PendingOrder order = pendingOrders.get(tradingRecord);
        if (order == null) {
            return;
        }
        expireOrder(lastProcessedIndex, tradingRecord, order);
    }

    private void expireIfStale(int index, TradingRecord tradingRecord) {
        PendingOrder order = pendingOrders.get(tradingRecord);
        if (order == null || index <= order.expiryIndex) {
            return;
        }
        expireOrder(index, tradingRecord, order);
    }

    private void expireOrder(int index, TradingRecord tradingRecord, PendingOrder order) {
        if (shouldCommitPartial(order, tradingRecord)) {
            tradingRecord.operate(order.toTrade(tradingRecord));
        }
        addRejectedOrder(tradingRecord, order.toExpiryRejection(index));
        pendingOrders.remove(tradingRecord);
    }

    private static boolean shouldCommitPartial(PendingOrder order, TradingRecord tradingRecord) {
        if (!order.hasAnyFill()) {
            return false;
        }
        if (order.tradeType == tradingRecord.getStartingType()) {
            return true;
        }
        return !tradingRecord.getOpenPositions().isEmpty();
    }

    private static void validateRatio(Num ratio, String name) {
        Objects.requireNonNull(ratio, name);
        if (ratio.isNaN() || ratio.isNegative()) {
            throw new IllegalArgumentException(name + " must be positive or zero");
        }
    }

    private int resolveActivationIndex(int referenceIndex) {
        if (priceSource == PriceSource.CURRENT_CLOSE) {
            return referenceIndex + 1;
        }
        return referenceIndex;
    }

    private Num toStopPrice(Num reference, TradeType tradeType) {
        Num one = reference.getNumFactory().one();
        if (tradeType == TradeType.BUY) {
            return reference.multipliedBy(one.plus(stopTriggerRatio));
        }
        return reference.multipliedBy(one.minus(stopTriggerRatio));
    }

    private Num toLimitPrice(Num reference, TradeType tradeType) {
        Num one = reference.getNumFactory().one();
        if (tradeType == TradeType.BUY) {
            return reference.multipliedBy(one.plus(limitOffsetRatio));
        }
        return reference.multipliedBy(one.minus(limitOffsetRatio));
    }

    private Num fillAmount(Num remainingAmount, Num barVolume) {
        Num availableAmount = remainingAmount;
        if (!Num.isNaNOrNull(barVolume)) {
            if (!barVolume.isPositive()) {
                return remainingAmount.getNumFactory().zero();
            }
            availableAmount = barVolume.multipliedBy(maxBarParticipationRate);
        }
        if (availableAmount.isNaN() || availableAmount.isNegativeOrZero()) {
            return remainingAmount.getNumFactory().zero();
        }
        if (availableAmount.isGreaterThan(remainingAmount)) {
            return remainingAmount;
        }
        return availableAmount;
    }

    private static boolean triggerReached(TradeType tradeType, Bar bar, Num stopPrice) {
        if (tradeType == TradeType.BUY) {
            return bar.getHighPrice().isGreaterThanOrEqual(stopPrice);
        }
        return bar.getLowPrice().isLessThanOrEqual(stopPrice);
    }

    private static boolean limitReachable(TradeType tradeType, Bar bar, Num limitPrice) {
        if (tradeType == TradeType.BUY) {
            return bar.getLowPrice().isLessThanOrEqual(limitPrice);
        }
        return bar.getHighPrice().isGreaterThanOrEqual(limitPrice);
    }

    private static Num amountOrZero(Num amount, BarSeries barSeries) {
        if (amount == null || amount.isNaN()) {
            return barSeries.numFactory().zero();
        }
        return amount;
    }

    private static Num resolveRequestedAmount(TradingRecord tradingRecord, Num defaultAmount) {
        if (tradingRecord.isClosed()) {
            return defaultAmount;
        }
        Position currentPosition = tradingRecord.getCurrentPosition();
        if (currentPosition.isOpened() && currentPosition.getEntry() != null
                && currentPosition.getEntry().getAmount() != null && !currentPosition.getEntry().getAmount().isNaN()) {
            return currentPosition.getEntry().getAmount();
        }
        return defaultAmount;
    }

    private void addRejectedOrder(TradingRecord tradingRecord, RejectedOrder rejection) {
        rejectedOrders.computeIfAbsent(tradingRecord, ignored -> new ArrayList<>()).add(rejection);
    }

    /**
     * Rejected stop-limit order metadata.
     *
     * @param signalIndex     strategy signal index
     * @param rejectionIndex  bar index where rejection happened
     * @param tradeType       trade side
     * @param requestedAmount requested amount
     * @param filledAmount    amount filled before rejection
     * @param reason          rejection reason
     * @since 0.22.4
     */
    public record RejectedOrder(int signalIndex, int rejectionIndex, TradeType tradeType, Num requestedAmount,
            Num filledAmount, String reason) {
    }

    /**
     * Snapshot of a pending stop-limit order.
     *
     * @param signalIndex     strategy signal index
     * @param activationIndex first index where order can execute
     * @param tradeType       order side
     * @param requestedAmount requested amount
     * @param filledAmount    filled amount
     * @param stopPrice       stop trigger price
     * @param limitPrice      limit price
     * @param expiryIndex     last fillable bar index
     * @param triggered       true if stop trigger was reached
     * @param fills           current fills
     * @since 0.22.4
     */
    public record PendingOrderSnapshot(int signalIndex, int activationIndex, TradeType tradeType, Num requestedAmount,
            Num filledAmount, Num stopPrice, Num limitPrice, int expiryIndex, boolean triggered,
            List<TradeFill> fills) {
    }

    private static final class PendingOrder {
        private final int signalIndex;
        private final int activationIndex;
        private final TradeType tradeType;
        private final Num requestedAmount;
        private final Num stopPrice;
        private final Num limitPrice;
        private final int expiryIndex;
        private boolean triggered;
        private Num filledAmount;
        private final List<TradeFill> fills;

        private PendingOrder(int signalIndex, int activationIndex, TradeType tradeType, Num requestedAmount,
                Num stopPrice, Num limitPrice, int expiryIndex) {
            this.signalIndex = signalIndex;
            this.activationIndex = activationIndex;
            this.tradeType = tradeType;
            this.requestedAmount = requestedAmount;
            this.stopPrice = stopPrice;
            this.limitPrice = limitPrice;
            this.expiryIndex = expiryIndex;
            this.triggered = false;
            this.filledAmount = requestedAmount.getNumFactory().zero();
            this.fills = new ArrayList<>();
        }

        private Num remainingAmount() {
            return requestedAmount.minus(filledAmount);
        }

        private void recordFill(int index, Bar bar, Num price, Num amount) {
            fills.add(new TradeFill(index, bar.getEndTime(), price, amount, sideOf(tradeType)));
            filledAmount = filledAmount.plus(amount);
        }

        private boolean isCompletelyFilled() {
            return !requestedAmount.minus(filledAmount).isPositive();
        }

        private boolean hasAnyFill() {
            return filledAmount.isPositive();
        }

        private Trade toTrade(TradingRecord tradingRecord) {
            return Trade.fromFills(tradeType, fills, tradingRecord.getTransactionCostModel());
        }

        private RejectedOrder toExpiryRejection(int rejectionIndex) {
            return new RejectedOrder(signalIndex, rejectionIndex, tradeType, requestedAmount, filledAmount,
                    "Stop-limit order expired before filling requested amount");
        }

        private PendingOrderSnapshot snapshot() {
            return new PendingOrderSnapshot(signalIndex, activationIndex, tradeType, requestedAmount, filledAmount,
                    stopPrice, limitPrice, expiryIndex, triggered, List.copyOf(fills));
        }

        private static ExecutionSide sideOf(TradeType tradeType) {
            if (tradeType == TradeType.BUY) {
                return ExecutionSide.BUY;
            }
            return ExecutionSide.SELL;
        }
    }
}
