agent/main.py

"""
Agent Orchestrator — 6-Stage Pipeline: Observe → Reason → Plan → Authorize → Execute → Verify
===============================================================================================
Main entry point for the Dynamic RWA Yield Router. Orchestrates the full
autonomous agent loop with configurable intervals and safety guardrails.
"""

import asyncio
import json
import logging
import os
import signal
import sys
import time
from datetime import datetime, timezone
from typing import Dict, List, Optional

import numpy as np

from config.constants import (
    MANTLE_CHAIN_ID,
    MANTLE_RPC_URL,
    PortfolioConfig,
    RiskConfig,
    RLConfig,
    AGENT_METADATA,
    LOG_FORMAT,
    LOG_DATE_FORMAT,
    LOG_LEVEL,
)
from agent.data_pipeline import DataPipeline, YieldSnapshot
from agent.rl_optimizer import PPOYieldOptimizer, RWAYieldEnv, Backtester
from agent.risk_manager import RiskManager, RiskLevel
from agent.executor import OnChainExecutor, RebalancePlan
from agent.strategy_reporter import StrategyReporter

logger = logging.getLogger("orchestrator")


class AgentState:
    """Mutable agent state tracking."""
    
    def __init__(self, initial_capital: float = 100_000.0):
        self.portfolio_value: float = initial_capital
        self.current_weights: np.ndarray = np.array([0.40, 0.35, 0.25])  # USDY, mETH, MI4
        self.target_weights: Optional[np.ndarray] = None
        
        self.last_rebalance_time: float = 0.0
        self.total_rebalances: int = 0
        self.total_gas_spent_usd: float = 0.0
        
        self.pending_plans: List[RebalancePlan] = []
        self.executed_plans: List[RebalancePlan] = []
        
        # Performance tracking
        self.initial_capital = initial_capital
        self.peak_value = initial_capital
        self.cumulative_yield_pct: float = 0.0
        self.start_time: float = time.time()
    
    def to_dict(self) -> Dict:
        return {
            "portfolio_value": self.portfolio_value,
            "current_weights": {
                "USDY": float(self.current_weights[0]),
                "mETH": float(self.current_weights[1]),
                "MI4": float(self.current_weights[2]),
            },
            "total_return_pct": (self.portfolio_value / self.initial_capital - 1) * 100,
            "total_rebalances": self.total_rebalances,
            "total_gas_spent": self.total_gas_spent_usd,
            "uptime_hours": (time.time() - self.start_time) / 3600,
        }


class YieldRouterAgent:
    """
    Autonomous AI agent for RWA yield optimization on Mantle L2.
    
    Pipeline Stages:
    1. OBSERVE  — Fetch real-time yield, price, and macro data
    2. REASON   — RL policy predicts optimal allocation weights
    3. PLAN     — Compute rebalancing trades if weights deviate
    4. AUTHORIZE — Risk manager validates the plan
    5. EXECUTE  — Construct unsigned transactions
    6. VERIFY   — Confirm execution and update state
    
    Safety:
    - Circuit breaker halts trading on excessive drawdown or depeg
    - Position limits enforce diversification (5-60% per asset)
    - Rebalancing cooldown prevents excessive trading
    - All transactions are unsigned — requires external signing
    """
    
    def __init__(
        self,
        wallet_address: str = "0x0000000000000000000000000000000000000000",
        initial_capital: float = 100_000.0,
        portfolio_config: Optional[PortfolioConfig] = None,
        risk_config: Optional[RiskConfig] = None,
        rl_config: Optional[RLConfig] = None,
        rpc_url: str = MANTLE_RPC_URL,
        model_path: Optional[str] = None,
    ):
        self.wallet = wallet_address
        self.portfolio_cfg = portfolio_config or PortfolioConfig()
        self.risk_cfg = risk_config or RiskConfig()
        self.rl_cfg = rl_config or RLConfig()
        
        # Initialize components
        self.data_pipeline = DataPipeline(rpc_url=rpc_url)
        self.rl_optimizer = PPOYieldOptimizer(
            model_path=model_path or "models/ppo_yield_router",
            learning_rate=self.rl_cfg.learning_rate,
            total_timesteps=self.rl_cfg.total_timesteps,
        )
        self.risk_manager = RiskManager(
            risk_config=self.risk_cfg,
            portfolio_config=self.portfolio_cfg,
        )
        self.executor = OnChainExecutor(
            wallet_address=wallet_address,
            portfolio_config=self.portfolio_cfg,
        )
        self.strategy_reporter = StrategyReporter(use_llm=bool(os.getenv("OPENAI_API_KEY")))
        
        # State
        self.state = AgentState(initial_capital=initial_capital)
        
        # Latest market data
        self.latest_snapshot: Optional[YieldSnapshot] = None
        
        # Running flag
        self._running = False
        
        logger.info(
            f"YieldRouterAgent initialized | wallet={wallet_address[:10]}... | "
            f"capital=${initial_capital:,.0f} | chain={MANTLE_CHAIN_ID}"
        )
    
    # ────────────────────── Stage 1: OBSERVE ──────────────────────
    
    async def observe(self) -> YieldSnapshot:
        """
        Stage 1: Fetch all market data and construct observation state.
        
        Data sources: DeFiLlama, CoinGecko/Bybit, Mantle RPC, FRED
        Output: YieldSnapshot with normalized state vector
        """
        logger.info("📡 Stage 1: OBSERVE — fetching market data...")
        
        snapshot = await self.data_pipeline.get_snapshot()
        self.latest_snapshot = snapshot
        
        logger.info(
            f"  Yields: USDY={snapshot.usdy_apy:.2f}% mETH={snapshot.meth_apy:.2f}% MI4={snapshot.mi4_apy:.2f}%"
        )
        logger.info(
            f"  Prices: ETH=${snapshot.eth_price:.0f} BTC=${snapshot.btc_price:.0f} MNT=${snapshot.mnt_price:.3f}"
        )
        logger.info(
            f"  Health: USDY_peg={snapshot.usdy_peg:.4f} mETH_peg={snapshot.meth_peg:.4f}"
        )
        
        return snapshot
    
    # ────────────────────── Stage 2: REASON ──────────────────────
    
    def reason(self, snapshot: YieldSnapshot) -> np.ndarray:
        """
        Stage 2: RL policy predicts optimal portfolio weights.
        
        Input: Normalized state vector from snapshot + current weights
        Output: Target weights [USDY, mETH, MI4]
        """
        logger.info("🧠 Stage 2: REASON — RL policy inference...")
        
        # Construct full state: market features + current portfolio weights
        market_state = snapshot.to_state_vector()
        portfolio_state = self.state.current_weights
        full_state = np.concatenate([market_state, portfolio_state])
        
        # Predict optimal weights
        target_weights = self.rl_optimizer.predict(full_state)
        
        logger.info(
            f"  Target weights: USDY={target_weights[0]:.3f} "
            f"mETH={target_weights[1]:.3f} MI4={target_weights[2]:.3f}"
        )
        
        return target_weights
    
    # ────────────────────── Stage 3: PLAN ──────────────────────
    
    def plan(
        self, target_weights: np.ndarray, snapshot: YieldSnapshot
    ) -> Optional[RebalancePlan]:
        """
        Stage 3: Compute rebalancing plan if weights deviate enough.
        
        Checks:
        - Weight drift exceeds rebalance threshold
        - Minimum time since last rebalance
        - Gas cost vs. expected benefit
        
        Output: RebalancePlan or None (if no rebalance needed)
        """
        logger.info("📋 Stage 3: PLAN — computing rebalancing trades...")
        
        # Check rebalance cooldown
        hours_since_rebalance = (time.time() - self.state.last_rebalance_time) / 3600
        if hours_since_rebalance < self.portfolio_cfg.min_rebalance_interval_hours:
            logger.info(
                f"  ⏱️ Cooldown: {hours_since_rebalance:.1f}h since last rebalance "
                f"(min: {self.portfolio_cfg.min_rebalance_interval_hours}h)"
            )
            return None
        
        # Check weight drift
        drift = np.sum(np.abs(target_weights - self.state.current_weights))
        if drift < self.portfolio_cfg.rebalance_threshold_pct:
            logger.info(f"  ✅ Weights within threshold (drift={drift:.4f})")
            return None
        
        # Estimate asset prices (for trade sizing)
        asset_prices = {
            "USDY": 1.0,  # pegged to USD
            "mETH": snapshot.eth_price * snapshot.meth_peg,
            "MI4": 100.0,  # approximate NAV
        }
        
        # Build rebalance plan
        plan = self.executor.build_rebalance_plan(
            current_weights=self.state.current_weights.tolist(),
            target_weights=target_weights.tolist(),
            portfolio_value_usd=self.state.portfolio_value,
            asset_prices=asset_prices,
        )
        
        logger.info(f"  Plan: {plan.human_summary}")
        return plan
    
    # ────────────────────── Stage 4: AUTHORIZE ──────────────────────
    
    def authorize(
        self, plan: RebalancePlan, snapshot: YieldSnapshot
    ) -> bool:
        """
        Stage 4: Risk manager validates the rebalancing plan.
        
        Checks:
        - Depeg detection
        - Volatility regime
        - Position limits
        - Drawdown protection
        - Circuit breaker state
        
        Output: True if plan is approved, False if blocked
        """
        logger.info("🛡️ Stage 4: AUTHORIZE — risk assessment...")
        
        target_weights = np.array(plan.target_weights)
        
        assessment = self.risk_manager.assess_risk(
            proposed_weights=target_weights,
            current_weights=self.state.current_weights,
            snapshot=snapshot,
            portfolio_value=self.state.portfolio_value,
        )
        
        if assessment.warnings:
            for warning in assessment.warnings:
                logger.warning(f"  {warning}")
        
        if assessment.emergency_exit_recommended:
            logger.critical("  🚨 EMERGENCY EXIT recommended!")
            # Override with safe-haven weights
            safe_weights = self.risk_manager.get_emergency_exit_weights()
            self.state.target_weights = safe_weights
        
        if assessment.adjusted_weights is not None and not np.allclose(
            assessment.adjusted_weights, target_weights
        ):
            logger.info(
                f"  Weights adjusted by risk manager: "
                f"{assessment.adjusted_weights.tolist()}"
            )
            self.state.target_weights = assessment.adjusted_weights
        
        logger.info(
            f"  Risk: {assessment.overall_risk.value} "
            f"(score={assessment.risk_score:.3f}) "
            f"Approved: {assessment.rebalance_approved}"
        )
        
        return assessment.rebalance_approved
    
    # ────────────────────── Stage 5: EXECUTE ──────────────────────
    
    def execute(self, plan: RebalancePlan) -> Dict:
        """
        Stage 5: Output unsigned transactions for external signing.
        
        In production, these transactions would be submitted to an
        ERC-4337 bundler or multisig for execution. This agent only
        constructs the payloads — it never holds private keys.
        
        Output: Dict with all transaction payloads
        """
        logger.info("⚡ Stage 5: EXECUTE — constructing transactions...")
        
        result = {
            "plan": plan.to_dict(),
            "approvals": [tx.to_dict() for tx in plan.approvals],
            "trades": [tx.to_dict() for tx in plan.trades],
            "total_transactions": len(plan.approvals) + len(plan.trades),
            "estimated_gas_usd": plan.estimated_gas_usd,
        }
        
        logger.info(
            f"  Transactions ready: {result['total_transactions']} "
            f"({len(plan.approvals)} approvals + {len(plan.trades)} swaps)"
        )
        
        # In simulation mode, we update state directly
        self.state.current_weights = np.array(plan.target_weights)
        self.state.last_rebalance_time = time.time()
        self.state.total_rebalances += 1
        self.state.total_gas_spent_usd += plan.estimated_gas_usd
        self.state.executed_plans.append(plan)
        
        return result
    
    # ────────────────────── Stage 6: VERIFY ──────────────────────
    
    async def verify(self, execution_result: Dict) -> bool:
        """
        Stage 6: Verify execution and update portfolio state.
        
        In production, this would check on-chain transaction receipts.
        In simulation, we verify state consistency.
        """
        logger.info("✅ Stage 6: VERIFY — confirming execution...")
        
        # Verify weights sum to 1
        weight_sum = self.state.current_weights.sum()
        if abs(weight_sum - 1.0) > 0.001:
            logger.error(f"  ❌ Weight sum error: {weight_sum:.4f}")
            self.state.current_weights /= weight_sum
            return False
        
        # Verify no negative weights
        if np.any(self.state.current_weights < 0):
            logger.error("  ❌ Negative weights detected!")
            return False
        
        # Update portfolio value (simulate yield accrual)
        if self.latest_snapshot:
            yields = np.array([
                self.latest_snapshot.usdy_apy,
                self.latest_snapshot.meth_apy,
                self.latest_snapshot.mi4_apy,
            ])
            # Hourly yield accrual
            hourly_yield = np.dot(self.state.current_weights, yields / 100 / 365 / 24)
            self.state.portfolio_value *= (1 + hourly_yield)
            self.state.peak_value = max(self.state.peak_value, self.state.portfolio_value)
        
        logger.info(
            f"  Portfolio: ${self.state.portfolio_value:,.2f} | "
            f"Weights: {self.state.current_weights.tolist()} | "
            f"Rebalances: {self.state.total_rebalances}"
        )
        
        return True
    
    # ────────────────────── Full Pipeline ──────────────────────
    
    async def run_cycle(self) -> Dict:
        """
        Execute one full observation-to-verification cycle.
        
        Returns a summary dict of the cycle results.
        """
        cycle_start = time.time()
        
        try:
            # Stage 1: Observe
            snapshot = await self.observe()
            
            # Stage 2: Reason
            target_weights = self.reason(snapshot)
            
            # Stage 3: Plan
            plan = self.plan(target_weights, snapshot)
            
            if plan is None:
                logger.info("🔄 Cycle complete: no rebalance needed")
                # Still verify state
                await self.verify({})
                return {
                    "status": "no_rebalance",
                    "cycle_time_s": time.time() - cycle_start,
                    "state": self.state.to_dict(),
                }
            
            # Stage 4: Authorize
            approved = self.authorize(plan, snapshot)
            
            if not approved:
                logger.warning("🚫 Cycle complete: rebalance blocked by risk manager")
                return {
                    "status": "blocked",
                    "cycle_time_s": time.time() - cycle_start,
                    "state": self.state.to_dict(),
                    "risk": self.risk_manager.get_risk_summary(),
                }
            
            # Stage 5: Execute
            execution_result = self.execute(plan)
            
            # Stage 6: Verify
            verified = await self.verify(execution_result)
            
            status = "success" if verified else "verification_failed"
            logger.info(f"🏁 Cycle complete: {status} in {time.time() - cycle_start:.1f}s")
            
            return {
                "status": status,
                "cycle_time_s": time.time() - cycle_start,
                "execution": execution_result,
                "state": self.state.to_dict(),
            }
            
        except Exception as e:
            logger.error(f"❌ Cycle error: {e}", exc_info=True)
            return {
                "status": "error",
                "error": str(e),
                "cycle_time_s": time.time() - cycle_start,
                "state": self.state.to_dict(),
            }
    
    # ────────────────────── Continuous Operation ──────────────────
    
    async def run(
        self,
        interval_seconds: int = 3600,  # default: every hour
        max_cycles: Optional[int] = None,
        generate_reports: bool = True,
        report_interval_cycles: int = 168,  # weekly at hourly intervals
    ):
        """
        Run the agent in continuous loop.
        
        Args:
            interval_seconds: Time between cycles
            max_cycles: Stop after N cycles (None = run forever)
            generate_reports: Whether to generate strategy reports
            report_interval_cycles: How often to generate reports
        """
        self._running = True
        cycle_count = 0
        
        logger.info(
            f"🚀 Agent starting | interval={interval_seconds}s | "
            f"max_cycles={'∞' if max_cycles is None else max_cycles}"
        )
        
        while self._running:
            cycle_count += 1
            
            if max_cycles and cycle_count > max_cycles:
                logger.info(f"Reached max cycles ({max_cycles}), stopping.")
                break
            
            logger.info(f"\n{'='*60}")
            logger.info(f"Cycle #{cycle_count} | {datetime.now(timezone.utc).isoformat()}")
            logger.info(f"{'='*60}")
            
            result = await self.run_cycle()
            
            # Generate strategy report periodically
            if generate_reports and cycle_count % report_interval_cycles == 0:
                await self._generate_report()
            
            # Log state
            state = self.state.to_dict()
            logger.info(
                f"State: value=${state['portfolio_value']:,.2f} | "
                f"return={state['total_return_pct']:+.2f}% | "
                f"rebalances={state['total_rebalances']} | "
                f"gas=${state['total_gas_spent']:,.4f}"
            )
            
            # Save state checkpoint
            self._save_checkpoint()
            
            if self._running and (max_cycles is None or cycle_count < max_cycles):
                logger.info(f"💤 Sleeping {interval_seconds}s until next cycle...")
                await asyncio.sleep(interval_seconds)
        
        logger.info("🛑 Agent stopped.")
        await self.data_pipeline.close()
    
    async def _generate_report(self):
        """Generate and log a strategy report."""
        try:
            report = self.strategy_reporter.generate_report(
                current_weights={
                    "USDY": float(self.state.current_weights[0]),
                    "mETH": float(self.state.current_weights[1]),
                    "MI4": float(self.state.current_weights[2]),
                },
                portfolio_value=self.state.portfolio_value,
                period_return=(self.state.portfolio_value / self.state.initial_capital - 1) * 100,
                yield_data={
                    "usdy": self.latest_snapshot.usdy_apy if self.latest_snapshot else 4.25,
                    "meth": self.latest_snapshot.meth_apy if self.latest_snapshot else 3.5,
                    "mi4": self.latest_snapshot.mi4_apy if self.latest_snapshot else 5.0,
                },
                risk_summary=self.risk_manager.get_risk_summary(),
                execution_history=self.executor.get_execution_history(),
            )
            
            logger.info(f"\n📊 Strategy Report Generated:\n{report.full_report[:500]}...")
            
            # Save report
            os.makedirs("reports", exist_ok=True)
            with open(f"reports/{report.report_id}.json", "w") as f:
                json.dump(report.to_dict(), f, indent=2, default=str)
                
        except Exception as e:
            logger.error(f"Report generation failed: {e}")
    
    def _save_checkpoint(self):
        """Save agent state checkpoint."""
        os.makedirs("checkpoints", exist_ok=True)
        checkpoint = {
            "timestamp": time.time(),
            "state": self.state.to_dict(),
            "risk": self.risk_manager.get_risk_summary(),
            "agent_metadata": AGENT_METADATA,
        }
        with open("checkpoints/latest.json", "w") as f:
            json.dump(checkpoint, f, indent=2, default=str)
    
    def stop(self):
        """Signal the agent to stop after current cycle."""
        self._running = False
        logger.info("Stop signal received.")
    
    # ────────────────────── Training ──────────────────────
    
    def train_rl_agent(self, total_timesteps: Optional[int] = None):
        """Train or retrain the RL policy."""
        logger.info("🎓 Training RL agent...")
        self.rl_optimizer.train(total_timesteps=total_timesteps)
        logger.info("✅ RL training complete.")
    
    def backtest(self, n_episodes: int = 10) -> Dict:
        """Run backtesting simulation."""
        logger.info(f"📊 Running backtest ({n_episodes} episodes)...")
        env = RWAYieldEnv()
        backtester = Backtester(self.rl_optimizer, env)
        results = backtester.run_backtest(n_episodes)
        
        if results.get("rl_agent"):
            avg_return = np.mean([r["total_return"] for r in results["rl_agent"]])
            avg_sharpe = np.mean([r["sharpe"] for r in results["rl_agent"]])
            avg_dd = np.mean([r["max_drawdown"] for r in results["rl_agent"]])
            
            logger.info(
                f"Backtest results: "
                f"Avg Return={avg_return:.2f}% | "
                f"Avg Sharpe={avg_sharpe:.2f} | "
                f"Avg MaxDD={avg_dd:.4f}"
            )
        
        return results


# ─────────────────────── CLI Entry Point ────────────────────────────

def setup_logging():
    """Configure logging."""
    logging.basicConfig(
        level=getattr(logging, LOG_LEVEL),
        format=LOG_FORMAT,
        datefmt=LOG_DATE_FORMAT,
        handlers=[
            logging.StreamHandler(sys.stdout),
            logging.FileHandler("agent.log", mode="a"),
        ],
    )


async def main():
    """CLI entry point."""
    import argparse
    
    parser = argparse.ArgumentParser(description="Dynamic RWA Yield Router Agent")
    parser.add_argument("--mode", choices=["run", "train", "backtest", "demo"], default="demo")
    parser.add_argument("--wallet", type=str, default="0x" + "0" * 40)
    parser.add_argument("--capital", type=float, default=100_000.0)
    parser.add_argument("--interval", type=int, default=3600, help="Seconds between cycles")
    parser.add_argument("--cycles", type=int, default=None, help="Max cycles (None=infinite)")
    parser.add_argument("--train-steps", type=int, default=100_000)
    parser.add_argument("--backtest-episodes", type=int, default=10)
    parser.add_argument("--rpc", type=str, default=MANTLE_RPC_URL)
    
    args = parser.parse_args()
    
    setup_logging()
    
    agent = YieldRouterAgent(
        wallet_address=args.wallet,
        initial_capital=args.capital,
        rpc_url=args.rpc,
    )
    
    # Handle shutdown signals
    def signal_handler(sig, frame):
        logger.info("Shutdown signal received...")
        agent.stop()
    
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    
    if args.mode == "train":
        agent.train_rl_agent(total_timesteps=args.train_steps)
        
    elif args.mode == "backtest":
        results = agent.backtest(n_episodes=args.backtest_episodes)
        print(json.dumps(results, indent=2, default=str))
        
    elif args.mode == "run":
        await agent.run(
            interval_seconds=args.interval,
            max_cycles=args.cycles,
        )
        
    elif args.mode == "demo":
        # Run 3 demo cycles with short interval
        logger.info("🎭 Running demo mode (3 cycles, 5s interval)...")
        await agent.run(
            interval_seconds=5,
            max_cycles=3,
            generate_reports=True,
            report_interval_cycles=3,
        )
        print("\n" + "=" * 60)
        print("Demo complete! Final state:")
        print(json.dumps(agent.state.to_dict(), indent=2))


if __name__ == "__main__":
    asyncio.run(main())