evotraders/backend/runtime/manager.py

from __future__ import annotations

import asyncio
import json
from datetime import datetime, UTC
from pathlib import Path
from typing import Any, Dict, List, Optional

from .agent_runtime import AgentRuntimeState
from .context import TradingRunContext
from .registry import RuntimeRegistry

_global_runtime_manager: Optional["TradingRuntimeManager"] = None
_shutdown_event: Optional[asyncio.Event] = None

# Lazy import to avoid circular dependency
_api_runtime = None


def _get_api_runtime():
    global _api_runtime
    if _api_runtime is None:
        from backend.api import runtime as api_runtime_module
        _api_runtime = api_runtime_module
    return _api_runtime


def set_global_runtime_manager(manager: "TradingRuntimeManager") -> None:
    global _global_runtime_manager
    _global_runtime_manager = manager
    # Sync to RuntimeState for consistency
    _get_api_runtime().register_runtime_manager(manager)


def clear_global_runtime_manager() -> None:
    global _global_runtime_manager
    _global_runtime_manager = None
    # Sync to RuntimeState for consistency
    _get_api_runtime().unregister_runtime_manager()


def get_global_runtime_manager() -> Optional["TradingRuntimeManager"]:
    return _global_runtime_manager


def set_shutdown_event(event: asyncio.Event) -> None:
    """Set the global shutdown event for signaling runtime stop."""
    global _shutdown_event
    _shutdown_event = event


def clear_shutdown_event() -> None:
    """Clear the global shutdown event."""
    global _shutdown_event
    _shutdown_event = None


def get_shutdown_event() -> Optional[asyncio.Event]:
    """Get the global shutdown event if set."""
    return _shutdown_event


def is_shutdown_requested() -> bool:
    """Check if shutdown has been requested."""
    return _shutdown_event is not None and _shutdown_event.is_set()


class TradingRuntimeManager:
    def __init__(self, config_name: str, run_dir: Path, bootstrap: Optional[Dict[str, Any]] = None) -> None:
        self.config_name = config_name
        self.run_dir = run_dir
        self.bootstrap = bootstrap or {}
        self.context: Optional[TradingRunContext] = None
        self.registry = RuntimeRegistry()
        self.current_session_key: Optional[str] = None
        self.events: List[Dict[str, Any]] = []
        self.pending_approvals: Dict[str, Dict[str, Any]] = {}
        self.snapshot_path = self.run_dir / "state" / "runtime_state.json"

    def prepare_run(self) -> TradingRunContext:
        self.run_dir.mkdir(parents=True, exist_ok=True)
        self.context = TradingRunContext(
            config_name=self.config_name,
            run_dir=self.run_dir,
            bootstrap_values=self.bootstrap,
        )
        self._persist_snapshot()
        return self.context

    def set_session_key(self, session_key: str) -> None:
        self.current_session_key = session_key
        self._persist_snapshot()

    def log_event(self, event: str, details: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        entry = {
            "timestamp": datetime.now(UTC).isoformat(),
            "event": event,
            "details": details or {},
            "session": self.current_session_key,
        }
        self.events.append(entry)
        self._persist_snapshot()
        return entry

    def register_agent(self, agent_id: str) -> AgentRuntimeState:
        state = AgentRuntimeState(agent_id=agent_id)
        self.registry.register(agent_id, state)
        self._persist_snapshot()
        return state

    def register_pending_approval(self, approval_id: str, payload: Dict[str, Any]) -> None:
        payload.setdefault("status", "pending")
        payload.setdefault("created_at", datetime.now(UTC).isoformat())
        self.pending_approvals[approval_id] = payload
        self._persist_snapshot()

    def update_agent_status(
        self,
        agent_id: str,
        status: str,
        session_key: Optional[str] = None,
    ) -> AgentRuntimeState:
        state = self.registry.get(agent_id)
        if state is None:
            state = self.register_agent(agent_id)
        effective_session = session_key or self.current_session_key
        state.update(status, effective_session)
        self._persist_snapshot()
        return state

    def get_agent_state(self, agent_id: str) -> Optional[AgentRuntimeState]:
        return self.registry.get(agent_id)

    def list_agents(self) -> list[str]:
        return self.registry.list_agents()

    def resolve_pending_approval(self, approval_id: str, resolved_by: str, status: str) -> None:
        entry = self.pending_approvals.get(approval_id)
        if not entry:
            return
        entry["status"] = status
        entry["resolved_at"] = datetime.now(UTC).isoformat()
        entry["resolved_by"] = resolved_by
        self._persist_snapshot()

    def list_pending_approvals(self) -> List[Dict[str, Any]]:
        return list(self.pending_approvals.values())

    def build_snapshot(self) -> Dict[str, Any]:
        return {
            "context": {
                "config_name": self.context.config_name,
                "run_dir": str(self.context.run_dir),
                "bootstrap_values": self.context.bootstrap_values,
            }
            if self.context
            else None,
            "current_session_key": self.current_session_key,
            "agents": [
                state.to_dict()
                for agent_id in self.registry.list_agents()
                if (state := self.registry.get(agent_id)) is not None
            ],
            "events": self.events,
            "pending_approvals": self.list_pending_approvals(),
        }

    def _persist_snapshot(self) -> None:
        self.snapshot_path.parent.mkdir(parents=True, exist_ok=True)
        self.snapshot_path.write_text(
            json.dumps(self.build_snapshot(), ensure_ascii=False, indent=2),
            encoding="utf-8",
        )