package com.quantai.trader.domain;

import static com.quantai.trader.util.TraderNumbers.*;

import com.quantai.trader.enums.PositionSide;

import java.math.BigDecimal;
import java.util.Map;
import java.util.Set;

public record RiskOutput(
        BigDecimal marketRiskProb,
        BigDecimal longPositionRiskProb,
        BigDecimal shortPositionRiskProb,
        BigDecimal marketPathRiskBps,
        BigDecimal longPositionPathRiskBps,
        BigDecimal shortPositionPathRiskBps,
        Map<String, BigDecimal> riskReasonScores
) {
    private static final Set<String> REQUIRED_REASON_KEYS = Set.of(
            "market_drawdown_prob", "volatility_expansion_prob", "spike_prob",
            "liquidity_deterioration_prob", "position_drawdown_prob");

    public RiskOutput {
        marketRiskProb = probability(marketRiskProb, "risk.marketRiskProb");
        longPositionRiskProb = probability(longPositionRiskProb, "risk.longPositionRiskProb");
        shortPositionRiskProb = probability(shortPositionRiskProb, "risk.shortPositionRiskProb");
        marketPathRiskBps = nonNegative(marketPathRiskBps, "risk.marketPathRiskBps");
        longPositionPathRiskBps = nonNegative(longPositionPathRiskBps, "risk.longPositionPathRiskBps");
        shortPositionPathRiskBps = nonNegative(shortPositionPathRiskBps, "risk.shortPositionPathRiskBps");
        riskReasonScores = checkedProbabilities(riskReasonScores, "risk.riskReasonScores");
    }

    public BigDecimal reasonScore(String key) {
        return riskReasonScores.get(requiredText(key, "risk reason key"));
    }

    public BigDecimal sideRiskProbFor(PositionSide side) {
        if (side == PositionSide.LONG) {
            return longPositionRiskProb;
        }
        if (side == PositionSide.SHORT) {
            return shortPositionRiskProb;
        }
        throw new IllegalArgumentException("position risk requires LONG or SHORT side");
    }

    public BigDecimal positionPathRiskBpsFor(PositionSide side) {
        if (side == PositionSide.LONG) {
            return longPositionPathRiskBps;
        }
        if (side == PositionSide.SHORT) {
            return shortPositionPathRiskBps;
        }
        throw new IllegalArgumentException("position path risk requires LONG or SHORT side");
    }

    private static Map<String, BigDecimal> checkedProbabilities(Map<String, BigDecimal> scores, String field) {
        Map<String, BigDecimal> source = scores == null ? Map.of() : scores;
        if (!source.keySet().containsAll(REQUIRED_REASON_KEYS)) {
            throw new IllegalArgumentException(field + " must contain " + REQUIRED_REASON_KEYS);
        }
        source.forEach((key, value) -> {
            requiredText(key, field + ".key");
            probability(value, field + "." + key);
        });
        return Map.copyOf(source);
    }
}