The-Ouroboros/src/strategy/models.py

"""Pydantic models for pre-market scenario planning.

Defines the data contracts for the proactive strategy system:
- AI generates DayPlaybook before market open (structured JSON scenarios)
- Local ScenarioEngine matches conditions during market hours (no API calls)
"""

from __future__ import annotations

from datetime import UTC, date, datetime
from enum import Enum

from pydantic import BaseModel, Field, field_validator


class ScenarioAction(str, Enum):
    """Actions that can be taken by scenarios."""

    BUY = "BUY"
    SELL = "SELL"
    HOLD = "HOLD"
    REDUCE_ALL = "REDUCE_ALL"


class MarketOutlook(str, Enum):
    """AI's assessment of market direction."""

    BULLISH = "bullish"
    NEUTRAL_TO_BULLISH = "neutral_to_bullish"
    NEUTRAL = "neutral"
    NEUTRAL_TO_BEARISH = "neutral_to_bearish"
    BEARISH = "bearish"


class PlaybookStatus(str, Enum):
    """Lifecycle status of a playbook."""

    PENDING = "pending"
    READY = "ready"
    FAILED = "failed"
    EXPIRED = "expired"


class StockCondition(BaseModel):
    """Condition fields for scenario matching (all optional, AND-combined).

    The ScenarioEngine evaluates all non-None fields as AND conditions.
    A condition matches only if ALL specified fields are satisfied.
    """

    rsi_below: float | None = None
    rsi_above: float | None = None
    volume_ratio_above: float | None = None
    volume_ratio_below: float | None = None
    price_above: float | None = None
    price_below: float | None = None
    price_change_pct_above: float | None = None
    price_change_pct_below: float | None = None

    def has_any_condition(self) -> bool:
        """Check if at least one condition field is set."""
        return any(
            v is not None
            for v in (
                self.rsi_below,
                self.rsi_above,
                self.volume_ratio_above,
                self.volume_ratio_below,
                self.price_above,
                self.price_below,
                self.price_change_pct_above,
                self.price_change_pct_below,
            )
        )


class StockScenario(BaseModel):
    """A single condition-action rule for one stock."""

    condition: StockCondition
    action: ScenarioAction
    confidence: int = Field(ge=0, le=100)
    allocation_pct: float = Field(ge=0, le=100, default=10.0)
    stop_loss_pct: float = Field(le=0, default=-2.0)
    take_profit_pct: float = Field(ge=0, default=3.0)
    rationale: str = ""


class StockPlaybook(BaseModel):
    """All scenarios for a single stock (ordered by priority)."""

    stock_code: str
    stock_name: str = ""
    scenarios: list[StockScenario] = Field(min_length=1)


class GlobalRule(BaseModel):
    """Portfolio-level rule (checked before stock-level scenarios)."""

    condition: str  # e.g. "portfolio_pnl_pct < -2.0"
    action: ScenarioAction
    rationale: str = ""


class CrossMarketContext(BaseModel):
    """Summary of another market's state for cross-market awareness."""

    market: str  # e.g. "US" or "KR"
    date: str
    total_pnl: float = 0.0
    win_rate: float = 0.0
    index_change_pct: float = 0.0  # e.g. KOSPI or S&P500 change
    key_events: list[str] = Field(default_factory=list)
    lessons: list[str] = Field(default_factory=list)


class DayPlaybook(BaseModel):
    """Complete playbook for a single trading day in a single market.

    Generated by PreMarketPlanner (1 Gemini call per market per day).
    Consumed by ScenarioEngine during market hours (0 API calls).
    """

    date: date
    market: str  # "KR" or "US"
    market_outlook: MarketOutlook = MarketOutlook.NEUTRAL
    generated_at: str = ""  # ISO timestamp
    gemini_model: str = ""
    token_count: int = 0
    global_rules: list[GlobalRule] = Field(default_factory=list)
    stock_playbooks: list[StockPlaybook] = Field(default_factory=list)
    default_action: ScenarioAction = ScenarioAction.HOLD
    context_summary: dict = Field(default_factory=dict)
    cross_market: CrossMarketContext | None = None

    @field_validator("stock_playbooks")
    @classmethod
    def validate_unique_stocks(cls, v: list[StockPlaybook]) -> list[StockPlaybook]:
        codes = [pb.stock_code for pb in v]
        if len(codes) != len(set(codes)):
            raise ValueError("Duplicate stock codes in playbook")
        return v

    def get_stock_playbook(self, stock_code: str) -> StockPlaybook | None:
        """Find the playbook for a specific stock."""
        for pb in self.stock_playbooks:
            if pb.stock_code == stock_code:
                return pb
        return None

    @property
    def scenario_count(self) -> int:
        """Total number of scenarios across all stocks."""
        return sum(len(pb.scenarios) for pb in self.stock_playbooks)

    @property
    def stock_count(self) -> int:
        """Number of stocks with scenarios."""
        return len(self.stock_playbooks)

    def model_post_init(self, __context: object) -> None:
        """Set generated_at if not provided."""
        if not self.generated_at:
            self.generated_at = datetime.now(UTC).isoformat()