from __future__ import annotations

from typing import ClassVar

from app.engine.indicators import supertrend as supertrend_series
from app.engine.strategies.base import Strategy, StrategyContext, register
from app.models import PositionSide


@register
class Supertrend(Strategy):
    """Trend following on the Supertrend indicator, long or short."""

    key: ClassVar[str] = "supertrend"
    name: ClassVar[str] = "Supertrend"
    description: ClassVar[str] = (
        "Go long when Supertrend flips to an uptrend and short when it flips to "
        "a downtrend. Uses the Supertrend line as a stop."
    )
    default_params: ClassVar[dict] = {
        "period": 10,
        "multiplier": 3.0,
        "allow_shorts": False,
        "size": 1.0,
    }
    param_help: ClassVar[dict] = {
        "period": "ATR period for the Supertrend bands",
        "multiplier": "ATR multiplier (band width)",
        "allow_shorts": "Also take short trades in downtrends",
        "size": "Fraction of equity per trade",
    }
    param_bounds: ClassVar[dict] = {
        "period": (2, 100),
        "multiplier": (0.5, 10.0),
        "size": (0.01, 1.0),
    }

    def on_candle(self, ctx: StrategyContext) -> None:
        period = int(self.params["period"])
        multiplier = float(self.params["multiplier"])
        if len(ctx.candles) < period + 2:
            return

        line_series, up_series = supertrend_series(ctx.candles, period, multiplier)
        line, up = line_series[-1], up_series[-1]
        prev_up = up_series[-2]
        if line is None or up is None or prev_up is None:
            return

        allow_shorts = bool(self.params.get("allow_shorts", False))
        size = float(self.params["size"])

        if ctx.side == PositionSide.LONG and not up:
            ctx.close(reason="Supertrend turned down")
            return
        if ctx.side == PositionSide.SHORT and up:
            ctx.close(reason="Supertrend turned up")
            return

        if ctx.is_flat and up and not prev_up:
            ctx.enter_long(size=size, stop_loss=line, reason="Supertrend flipped up")
        elif ctx.is_flat and allow_shorts and not up and prev_up:
            ctx.enter_short(size=size, stop_loss=line, reason="Supertrend flipped down")