crypto_trading/backend/app/engine/strategies/base.py

from __future__ import annotations

from abc import ABC, abstractmethod
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Any, ClassVar, Literal, Protocol

from app.engine.indicators import atr, ema, rolling_std, rsi, sma
from app.models import PositionSide
from app.schemas import Candle, ParamSpec, StrategyOut


class PositionView(Protocol):
    """Minimal position shape a strategy can see (DB row or backtest position)."""

    side: PositionSide
    qty: float
    entry_price: float


@dataclass(slots=True)
class Intent:
    """An order request emitted by a strategy. Executed by the engine."""

    action: Literal["enter_long", "enter_short", "close"]
    size: float | None = None
    stop_loss: float | None = None
    take_profit: float | None = None
    reason: str = ""


class StrategyContext:
    """Read-only market/portfolio view plus order intent helpers.

    A fresh context is created for every candle. Strategy code stays synchronous
    and side-effect free: it calls ``enter_long`` / ``enter_short`` / ``close``
    and the engine executes the resulting intent after ``on_candle`` returns.
    """

    def __init__(
        self,
        candles: Sequence[Candle],
        position: PositionView | None,
        equity: float,
        cash: float,
    ) -> None:
        self._candles = list(candles)
        self._position = position
        self._equity = equity
        self._cash = cash
        self._intents: list[Intent] = []

    # --- market state -------------------------------------------------
    @property
    def candles(self) -> list[Candle]:
        return self._candles

    @property
    def price(self) -> float:
        return self._candles[-1].close

    @property
    def closes(self) -> list[float]:
        return [c.close for c in self._candles]

    def high(self, lookback: int) -> float:
        return max(c.high for c in self._candles[-lookback:])

    def low(self, lookback: int) -> float:
        return min(c.low for c in self._candles[-lookback:])

    # --- portfolio state ----------------------------------------------
    @property
    def position(self) -> PositionView | None:
        return self._position

    @property
    def is_flat(self) -> bool:
        return self._position is None

    @property
    def side(self) -> PositionSide | None:
        return self._position.side if self._position else None

    @property
    def equity(self) -> float:
        return self._equity

    @property
    def cash(self) -> float:
        return self._cash

    # --- indicators ----------------------------------------------------
    def sma(self, period: int) -> float | None:
        return sma(self.closes, period)[-1]

    def ema(self, period: int) -> float | None:
        return ema(self.closes, period)[-1]

    def rsi(self, period: int = 14) -> float | None:
        return rsi(self.closes, period)[-1]

    def atr(self, period: int = 14) -> float | None:
        return atr(self._candles, period)[-1]

    def std(self, period: int = 20) -> float | None:
        return rolling_std(self.closes, period)[-1]

    # --- order intents ------------------------------------------------
    def enter_long(
        self,
        size: float | None = None,
        stop_loss: float | None = None,
        take_profit: float | None = None,
        reason: str = "",
    ) -> None:
        self._intents.append(Intent("enter_long", size, stop_loss, take_profit, reason))

    def enter_short(
        self,
        size: float | None = None,
        stop_loss: float | None = None,
        take_profit: float | None = None,
        reason: str = "",
    ) -> None:
        self._intents.append(Intent("enter_short", size, stop_loss, take_profit, reason))

    def close(self, reason: str = "") -> None:
        self._intents.append(Intent("close", reason=reason))

    def take_intents(self) -> list[Intent]:
        intents, self._intents = self._intents, []
        return intents


class Strategy(ABC):
    """Base class for trading strategies.

    Subclass it, set ``key``/``name``/``default_params`` and implement
    ``on_candle``. Save the file under ``app/engine/strategies/`` and it is
    auto-discovered.
    """

    key: ClassVar[str] = ""
    name: ClassVar[str] = ""
    description: ClassVar[str] = ""
    default_params: ClassVar[dict[str, Any]] = {}
    param_help: ClassVar[dict[str, str]] = {}
    param_bounds: ClassVar[dict[str, tuple[float, float]]] = {}

    def __init__(self, params: dict[str, Any] | None = None) -> None:
        unknown = set(params or {}) - set(self.default_params)
        if unknown:
            raise ValueError(f"unknown params for {self.key}: {sorted(unknown)}")
        self.params: dict[str, Any] = {**self.default_params, **(params or {})}

    @abstractmethod
    def on_candle(self, ctx: StrategyContext) -> None:
        raise NotImplementedError


def _infer_type(value: Any) -> Literal["int", "float", "bool", "str"]:
    if isinstance(value, bool):
        return "bool"
    if isinstance(value, int):
        return "int"
    if isinstance(value, float):
        return "float"
    return "str"


def param_specs(strategy_cls: type[Strategy]) -> list[ParamSpec]:
    specs: list[ParamSpec] = []
    for name, default in strategy_cls.default_params.items():
        bounds = strategy_cls.param_bounds.get(name)
        specs.append(
            ParamSpec(
                name=name,
                type=_infer_type(default),
                default=default,
                description=strategy_cls.param_help.get(name, ""),
                min=bounds[0] if bounds else None,
                max=bounds[1] if bounds else None,
            )
        )
    return specs


_REGISTRY: dict[str, type[Strategy]] = {}


def register(strategy_cls: type[Strategy]) -> type[Strategy]:
    """Class decorator that makes a strategy discoverable by key."""
    if not strategy_cls.key:
        raise ValueError(f"{strategy_cls.__name__} must define a non-empty 'key'")
    _REGISTRY[strategy_cls.key] = strategy_cls
    return strategy_cls


def get_strategy(key: str) -> type[Strategy]:
    try:
        return _REGISTRY[key]
    except KeyError:
        raise KeyError(f"unknown strategy '{key}'. available: {sorted(_REGISTRY)}") from None


def list_strategies() -> list[type[Strategy]]:
    return sorted(_REGISTRY.values(), key=lambda cls: cls.name or cls.key)


def strategy_spec(strategy_cls: type[Strategy]) -> StrategyOut:
    return StrategyOut(
        key=strategy_cls.key,
        name=strategy_cls.name or strategy_cls.key,
        description=strategy_cls.description,
        params=param_specs(strategy_cls),
    )