trex-core/trex_core/application/messaging_handlers.py

"""RPC handlers + event payload builders for messaging transport (REQ-021/022).

Handlers are transport-agnostic: they accept a parsed `RpcRequest` payload
and return an `RpcResponse`. Wiring them onto `MessagingTransport.consume_*`
happens in the runtime layer.
"""
from __future__ import annotations

import logging
from datetime import date as date_cls
from datetime import datetime, timezone
from typing import Optional
from uuid import UUID

from trex_core.application.intents import BuyIntent, IntentService, SellIntent
from trex_core.application.messaging import RpcRequest, RpcResponse
from trex_core.application.portfolio import (
    PauseStore,
    SnapshotStore,
    View,
    build_snapshot_envelope,
    slot_to_dict,
)
from trex_core.application.watchdog import WatchdogStore
from trex_core.domain.slot import (
    SlotNotFoundError,
    SlotRepository,
    SlotService,
)

_log = logging.getLogger(__name__)


# ---- intent RPC handler -----------------------------------------------------

def make_intent_rpc_handler(
    intent_service: IntentService,
    *,
    slot_service: Optional[SlotService] = None,
    watchdog_store: Optional[WatchdogStore] = None,
):
    """Build the RPC handler for `trex.intent.req`.

    Payload shape:
        {
          "intent_type": "buy" | "sell",
          "intent_id":   "<idempotency key, = correlation_id>",
          ... fields per BuyIntent / SellIntent ...
        }
      or (REQ-040):
        {"action": "update_thresholds", "slot_id": ...,
         "stop_loss_price"?, "take_profit_price"?, "forced_close_at"?}

    Response shape:
        {"intent_id": str, "status": "admitted"|"rejected", "slot_id": str|None, "reason": str|None}
      or for update_thresholds:
        {"status": "updated"|"ignored"|"error", "slot_id": str, "reason"?: str}

    Dedup uses `intent_log.try_insert(intent_id, ...)` (REQ-007 gate). The
    request's `correlation_id` is reused as `intent_id` so the messaging
    layer and the application share a single idempotency source of truth
    (REQ-023 AC-1).
    """

    def handler(req: RpcRequest) -> RpcResponse:
        body = req.payload or {}
        # REQ-040: update_thresholds action (separate from buy/sell intents).
        if str(body.get("action") or "").lower() == "update_thresholds":
            return _handle_update_thresholds(body, slot_service, watchdog_store)
        intent_id = str(body.get("intent_id") or req.correlation_id)
        intent_type = str(body.get("intent_type") or "").lower()
        try:
            if intent_type == "buy":
                # dcy-020 / REQ-061: buy is keyed by slot_id (canonical);
                # portfolio_id is the legacy dual-accept mapping. symbol is NOT
                # a client field — resolved from the universe. trading_date was
                # dropped from the contract at the cutover (BAT-030).
                buy = BuyIntent(
                    intent_id=intent_id,
                    portfolio_id=_opt_str(body.get("portfolio_id")) or None,
                    slot_id=(UUID(str(body["slot_id"])) if body.get("slot_id") else None),
                    qty=_opt_float(body.get("qty")),
                    allocated_capital=_opt_float(body.get("allocated_capital")),
                    entry_price=_opt_float(body.get("entry_price")),
                    strategy_id=_opt_str(body.get("strategy_id")),
                    sell_strategy_id=_opt_str(body.get("sell_strategy_id")),
                    forced_close_at=_opt_dt(body.get("forced_close_at")),
                    stop_loss_price=_opt_decimal(body.get("stop_loss_price")),
                    take_profit_price=_opt_decimal(body.get("take_profit_price")),
                )
                d = intent_service.submit_buy(buy)
            elif intent_type == "sell":
                sell = SellIntent(
                    intent_id=intent_id,
                    slot_id=UUID(str(body["slot_id"])),
                    reason=str(body.get("reason") or "discretionary_sell"),
                    sell_strategy_id=_opt_str(body.get("sell_strategy_id")),
                    # dcy-012 / REQ-007: optional sell limit price; absent => market sell.
                    limit_price=_opt_float(body.get("limit_price")),
                )
                d = intent_service.submit_sell(sell)
            else:
                return RpcResponse(payload={
                    "intent_id": intent_id, "status": "rejected",
                    "slot_id": None, "reason": f"unknown_intent_type:{intent_type}",
                })
        except (KeyError, ValueError) as exc:
            return RpcResponse(payload={
                "intent_id": intent_id, "status": "rejected",
                "slot_id": None, "reason": f"bad_request:{exc}",
            })
        return RpcResponse(payload={
            "intent_id": d.intent_id,
            "status": d.status,
            "slot_id": (str(d.slot_id) if d.slot_id else None),
            "reason": d.reason,
        })

    return handler


def _handle_update_thresholds(
    body: dict,
    slot_service: Optional[SlotService],
    watchdog_store: Optional[WatchdogStore],
) -> RpcResponse:
    """REQ-040: update stop_loss/take_profit/forced_close_at on a holding slot.

    Only the keys actually present in the payload are forwarded, so an absent
    field keeps its value while an explicit ``null`` clears it (AC-4 vs AC-5).
    On success the slot's watchdog last_seen_at is reset (AC-3) — receiving an
    update_thresholds is itself a strategy-module liveness signal.
    """
    if slot_service is None:
        return RpcResponse(payload={"status": "error", "reason": "update_thresholds_unsupported"})
    try:
        slot_id = UUID(str(body["slot_id"]))
    except (KeyError, ValueError) as exc:
        return RpcResponse(payload={"status": "error", "reason": f"bad_request:{exc}"})
    kwargs: dict = {}
    if "stop_loss_price" in body:
        kwargs["stop_loss_price"] = _opt_decimal(body["stop_loss_price"])
    if "take_profit_price" in body:
        kwargs["take_profit_price"] = _opt_decimal(body["take_profit_price"])
    if "forced_close_at" in body:
        kwargs["forced_close_at"] = _opt_dt(body["forced_close_at"])
    actor = str(body.get("actor") or "strategy")
    try:
        status = slot_service.update_thresholds(slot_id, actor=actor, **kwargs)
    except SlotNotFoundError:
        return RpcResponse(payload={
            "status": "ignored", "slot_id": str(slot_id), "reason": "slot_not_found",
        })
    if status == "updated" and watchdog_store is not None:
        watchdog_store.update(str(slot_id))  # AC-3: liveness signal
    payload = {"status": status, "slot_id": str(slot_id)}
    if status == "ignored":
        payload["reason"] = "not_holding"  # AC-6
    return RpcResponse(payload=payload)


# ---- snapshot RPC handler ---------------------------------------------------

def make_snapshot_rpc_handler(
    slot_repo: SlotRepository,
    slot_service: SlotService,
    *,
    instance_key: str = "default",
    snapshot_store: Optional[SnapshotStore] = None,
    pause_store: Optional[PauseStore] = None,
    max_active_slots: Optional[int] = None,
    watchdog_store: Optional[WatchdogStore] = None,
    watchdog_timeout_sec: int = 120,
    killswitch_store=None,
    ledger_outbox=None,
):
    """Build the RPC handler for `trex.snapshot.req` (REQ-028).

    When `snapshot_store` + `pause_store` are provided (production wiring),
    returns the BAT-008 envelope:
        {
          "assets":          list of {symbol, asset_id, cash_balance,
                                       asset_quantity, active_strategy_id,
                                       strategy_params},
          "view":            "all" | "buyable" | "held",
          "available_slots": int,
          "degraded":        bool,
          "reason":          DegradedReason value or null,
          "message":         sanitized text or null
        }

    Without those, falls back to the BAT-007 slot+capacity stub shape so
    older callers and existing tests don't break.
    """

    def handler(req: RpcRequest) -> RpcResponse:
        body = req.payload or {}
        day = _parse_date(body.get("trading_date")) if body.get("trading_date") else _today()

        # REQ-034: update watchdog for all holding slots visible in this RPC
        # (strategy polling == liveness signal). dcy-016/REQ-021: an operator
        # query (`touch_watchdog=false` / `actor="operator"`) must NOT count as
        # liveness — otherwise an operator looking at slots would suppress
        # safety-sell on a dead strategy (R-4). Default keeps legacy behavior.
        touch_watchdog = bool(body.get("touch_watchdog", True))
        if str(body.get("actor") or "").lower() == "operator":
            touch_watchdog = False
        if watchdog_store is not None and touch_watchdog:
            for s in slot_repo.list_holding_slots(instance_key):
                watchdog_store.update(str(s.slot_id))

        # --- BAT-008 envelope path (snapshot store wired) ---
        if snapshot_store is not None and pause_store is not None:
            requested_view = str(body.get("view") or "all").lower()
            try:
                view = View(requested_view)
            except ValueError:
                view = View.ALL
            strategy_key = _opt_str(body.get("strategy_key"))
            # dcy-020 / REQ-059: available = remaining admission headroom under
            # the M cap (max(0, M − active)); inactive count when no cap.
            cap = slot_service.compute_slot_capacity(max_active_slots)
            envelope = build_snapshot_envelope(
                store=snapshot_store,
                pause_store=pause_store,
                slot_repo=slot_repo,
                instance_key=instance_key,
                available_slots=cap.available_count,
                view=view,
                strategy_key=strategy_key,
                trading_date=day,
                killswitch_store=killswitch_store,
            )
            payload = envelope.to_dict()
            _attach_watchdog(payload, slot_repo, instance_key, watchdog_store, watchdog_timeout_sec)
            _apply_cash_estimate(payload, ledger_outbox)
            return RpcResponse(payload=payload)

        # --- BAT-007 legacy stub-shape path ---
        view_raw = str(body.get("view") or "all")
        strategy_filter = _opt_str(body.get("strategy_id"))
        slots = slot_repo.list_slots(instance_key)
        if view_raw == "open":
            slots = [s for s in slots if s.status.value in ("pending_buy", "holding", "pending_sell")]
        if strategy_filter is not None:
            slots = [s for s in slots if s.strategy_id == strategy_filter]
        cap = slot_service.compute_slot_capacity(max_active_slots)
        payload = {
            "view": view_raw,
            "strategy_id": strategy_filter,
            "slots": [_slot_to_dict(s) for s in slots],
            "capacity": {
                "inactive_count": cap.inactive_count,
                "holding_count": cap.holding_count,
                "pending_count": cap.pending_count,
                "active_count": cap.active_count,
                "available_count": cap.available_count,
                "total_slots": cap.total_slots,
                "scarcity_ratio": cap.scarcity_ratio,
            },
        }
        _attach_watchdog(payload, slot_repo, instance_key, watchdog_store, watchdog_timeout_sec)
        return RpcResponse(payload=payload)

    return handler


def _apply_cash_estimate(payload, ledger_outbox) -> None:
    """Overlay an estimated cash for assets with unsettled (pending) ledger
    writes (dcy-018 / REQ-029, OQ-4). cash_balance += pending cash delta and the
    asset is flagged `cash_estimated=true` until option-D write-back lands.
    Best-effort: any failure leaves the response unchanged (NFR-1)."""
    if ledger_outbox is None:
        return
    try:
        deltas = ledger_outbox.pending_cash_delta()
        if not deltas:
            return
        for a in payload.get("assets", []):
            d = deltas.get(a.get("asset_id"))
            if d:
                base = a.get("cash_balance") or 0.0
                a["cash_balance"] = base + d
                a["cash_estimated"] = True
    except Exception:  # noqa: BLE001 - estimate must never break the response
        return


def _attach_watchdog(payload, slot_repo, instance_key, watchdog_store, timeout_sec) -> None:
    """Attach holding-slot watchdog state to a snapshot response (dcy-017/REQ-034).

    `watchdog` = list of {slot_id, portfolio_id, symbol, expired, last_seen} for
    holding slots — lets `trex-core slots` surface watchdog ownership (strategy
    absence). Best-effort: any failure leaves the response unchanged (FR-4)."""
    if watchdog_store is None:
        return
    try:
        states = []
        for s in slot_repo.list_holding_slots(instance_key):
            sid = str(s.slot_id)
            last = watchdog_store.last_seen(sid)
            states.append({
                "slot_id": sid,
                "portfolio_id": s.portfolio_id,
                "symbol": s.symbol,
                "expired": watchdog_store.is_expired(sid, timeout_sec),
                "last_seen": last.isoformat() if last else None,
            })
        payload["watchdog"] = states
    except Exception:  # noqa: BLE001 - visibility must never break the response
        return


# ---- pause/resume RPC handler (REQ-032 AC-4) --------------------------------

def make_pause_rpc_handler(pause_store: PauseStore):
    """RPC handler for the `trex.pause.req` control queue.

    Payload: {"action": "pause" | "resume" | "list", "portfolio_id": "..."}
    Response: {"status": "ok"|"error", "paused": [list-of-portfolio_id], "reason"?: str}
    """

    def handler(req: RpcRequest) -> RpcResponse:
        body = req.payload or {}
        action = str(body.get("action") or "").lower()
        pid = _opt_str(body.get("portfolio_id"))
        actor = str(body.get("actor") or "operator")
        try:
            if action == "pause":
                if not pid:
                    return _err("portfolio_id required")
                pause_store.pause(pid, actor=actor)
            elif action == "resume":
                if not pid:
                    return _err("portfolio_id required")
                pause_store.resume(pid)
            elif action == "list":
                pass
            else:
                return _err(f"unknown action: {action!r}")
        except Exception as exc:  # noqa: BLE001
            return _err(f"pause store error: {exc.__class__.__name__}")
        return RpcResponse(payload={
            "status": "ok",
            "paused": sorted(pause_store.list_paused()),
            "action": action,
        })

    return handler


def _err(reason: str) -> RpcResponse:
    return RpcResponse(payload={"status": "error", "reason": reason})


# ---- control RPC handler (REQ-043, BAT-012) ---------------------------------

def make_control_rpc_handler(pause_store: PauseStore, *, kill_switch_service=None):
    """RPC handler for `trex.control.req` — action discriminator dispatch (REQ-043).

    Actions: `pause`/`resume`/`list` (REQ-032, delegated), `kill_switch` (REQ-045),
    `resume_kill_switch` (REQ-046). Unknown actions return an error response.
    """
    pause_handler = make_pause_rpc_handler(pause_store)

    def handler(req: RpcRequest) -> RpcResponse:
        body = req.payload or {}
        action = str(body.get("action") or "").lower()
        if action == "kill_switch":
            return _handle_kill_switch(body, kill_switch_service)
        if action == "resume_kill_switch":
            return _handle_resume_kill_switch(body, kill_switch_service)
        if action in ("pause", "resume", "list"):
            return pause_handler(req)
        return _err(f"unknown_action:{action}")

    return handler


def _killswitch_scope(body: dict):
    """Validate scope/portfolio_id. Returns (scope, pid) or (None, RpcResponse error)."""
    scope = str(body.get("scope") or "").lower()
    pid = _opt_str(body.get("portfolio_id"))
    if scope not in ("portfolio", "global"):
        return None, _err(f"bad_scope:{scope}")
    if scope == "portfolio" and not pid:
        return None, _err("portfolio_id_required")
    return (scope, pid), None


def _handle_kill_switch(body: dict, svc) -> RpcResponse:
    """REQ-045: liquidate holding slots in scope + day-exclude."""
    if svc is None:
        return _err("kill_switch_unsupported")
    parsed, err = _killswitch_scope(body)
    if err is not None:
        return err
    scope, pid = parsed
    actor = str(body.get("actor") or "operator")
    try:
        liquidated = svc.fire(scope=scope, portfolio_id=pid, actor=actor)
    except Exception as exc:  # noqa: BLE001
        return _err(f"kill_switch_failed:{exc.__class__.__name__}")
    payload = {"status": "ok", "scope": scope, "liquidated": liquidated}
    if pid:
        payload["portfolio_id"] = pid
    return RpcResponse(payload=payload)


def _handle_resume_kill_switch(body: dict, svc) -> RpcResponse:
    """REQ-046: clear the day-scoped exclusion (sold slots stay sold)."""
    if svc is None:
        return _err("resume_kill_switch_unsupported")
    parsed, err = _killswitch_scope(body)
    if err is not None:
        return err
    scope, pid = parsed
    try:
        svc.resume(scope=scope, portfolio_id=pid)
    except Exception as exc:  # noqa: BLE001
        return _err(f"resume_kill_switch_failed:{exc.__class__.__name__}")
    payload = {"status": "ok", "scope": scope}
    if pid:
        payload["portfolio_id"] = pid
    return RpcResponse(payload=payload)


# ---- event payload builder --------------------------------------------------

def slot_transition_event_payload(slot, *, event_type: str) -> dict:
    """Build the event payload for a slot transition (REQ-023 AC-5).

    Carries the *current state* + a monotonic sequence (`updated_at` iso) so
    receivers can apply last-wins without ordering guarantees.
    """
    return {
        "event_type": event_type,
        "slot_id": str(slot.slot_id),
        "strategy_id": slot.strategy_id,
        "symbol": slot.symbol,
        "status": slot.status.value,
        "qty": float(slot.qty) if slot.qty is not None else None,
        "avg_price": float(slot.avg_price) if slot.avg_price is not None else None,
        "portfolio_id": slot.portfolio_id,
        "sequence_at": (slot.updated_at or datetime.now(tz=timezone.utc)).isoformat(),
    }


# ---- helpers ----------------------------------------------------------------

def _opt_str(v) -> Optional[str]:
    if v is None or v == "":
        return None
    return str(v)


def _opt_decimal(v):
    if v is None:
        return None
    from decimal import Decimal
    return Decimal(str(v))


def _opt_float(v) -> Optional[float]:
    if v is None:
        return None
    return float(v)


def _opt_dt(v) -> Optional[datetime]:
    if v is None:
        return None
    if isinstance(v, datetime):
        return v
    return datetime.fromisoformat(str(v))


def _parse_date(v) -> date_cls:
    if isinstance(v, date_cls):
        return v
    return date_cls.fromisoformat(str(v))


def _today() -> date_cls:
    return datetime.now(tz=timezone.utc).date()


# BAT-022: canonical slot dict shared with build_snapshot_envelope (single
# source of truth so BAT-007 stub and BAT-008 envelope emit an identical
# `slots[]` shape). Local alias kept for existing call sites.
_slot_to_dict = slot_to_dict