"""Paper Trading module for OpenClaw trading system.

This module provides a PaperTradingManager that simulates real exchange behavior
without using real funds. It allows agents to test strategies safely with virtual
balances, simulated order execution, slippage, and fees.
"""

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

from pydantic import BaseModel, Field

from openclaw.exchange.base import Exchange, InsufficientFundsError
from openclaw.exchange.models import (
    Balance,
    Order,
    OrderSide,
    OrderStatus,
    OrderType,
    Position,
    Ticker,
)
from openclaw.exchange.mock import MockExchange
from openclaw.utils.logging import get_logger


class PaperTradeEntry(BaseModel):
    """Record of a paper trade execution.

    Attributes:
        timestamp: When the trade occurred
        symbol: Trading symbol
        side: Buy or sell
        order_type: Market or limit
        amount: Trade amount
        price: Execution price
        filled_amount: Actually filled amount
        fee: Trading fee charged
        slippage: Applied slippage percentage
        balance_before: Balance before trade
        balance_after: Balance after trade
        pnl: Realized profit/loss
    """

    timestamp: str = Field(..., description="ISO format timestamp")
    symbol: str = Field(..., description="Trading symbol")
    side: str = Field(..., description="Trade side")
    order_type: str = Field(..., description="Order type")
    amount: float = Field(..., gt=0, description="Order amount")
    price: float = Field(..., gt=0, description="Execution price")
    filled_amount: float = Field(..., ge=0, description="Filled amount")
    fee: float = Field(..., ge=0, description="Trading fee")
    slippage_pct: float = Field(..., description="Slippage percentage applied")
    balance_before: Dict[str, float] = Field(..., description="Balances before trade")
    balance_after: Dict[str, float] = Field(..., description="Balances after trade")
    pnl: float = Field(default=0.0, description="Realized PnL")
    order_id: str = Field(..., description="Order ID")
    metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional data")


class PaperTradingConfig(BaseModel):
    """Configuration for paper trading mode.

    Attributes:
        enabled: Whether paper trading is enabled
        initial_balances: Initial virtual balances {asset: amount}
        trading_fee_rate: Fee rate as decimal (default 0.001 = 0.1%)
        slippage_pct: Slippage percentage (default 0.05%)
        latency_ms: Simulated latency in milliseconds
        history_file: Path to trade history file
        price_source: Price source ('mock', 'ccxt', or 'websocket')
        enable_logging: Whether to log trades to file
    """

    enabled: bool = Field(default=True, description="Paper trading enabled")
    initial_balances: Dict[str, float] = Field(
        default_factory=lambda: {"USDT": 10000.0, "BTC": 0.1},
        description="Initial virtual balances",
    )
    trading_fee_rate: float = Field(
        default=0.001, ge=0, le=0.1, description="Trading fee rate"
    )
    slippage_pct: float = Field(
        default=0.05, ge=0, le=1.0, description="Slippage percentage"
    )
    latency_ms: float = Field(default=10.0, ge=0, description="Simulated latency")
    history_file: str = Field(
        default="logs/paper_trades.jsonl", description="Trade history file path"
    )
    price_source: str = Field(
        default="mock", description="Price source (mock/ccxt/websocket)"
    )
    enable_logging: bool = Field(default=True, description="Enable trade logging")


class PaperTradingManager:
    """Manager for paper trading simulation.

    Provides a complete paper trading environment that simulates real exchange
    behavior including:
    - Virtual balance management (USDT, BTC, etc.)
    - Order execution (market, limit orders)
    - Slippage and fee simulation
    - Trade history tracking
    - Integration with agents

    Args:
        config: Paper trading configuration
        exchange: Optional external exchange for price data
    """

    def __init__(
        self,
        config: Optional[PaperTradingConfig] = None,
        exchange: Optional[Exchange] = None,
    ):
        self.config = config or PaperTradingConfig()
        self.logger = get_logger("trading.paper_trading")

        # Create mock exchange for simulation
        self._exchange = MockExchange(
            name="paper",
            initial_balances=self.config.initial_balances.copy(),
            latency_ms=self.config.latency_ms,
            slippage_pct=self.config.slippage_pct,
        )

        # Override fee rate with config
        self._trading_fee_rate = self.config.trading_fee_rate

        # Trade history
        self._trade_history: List[PaperTradeEntry] = []
        self._positions: Dict[str, Position] = {}

        # Ensure log directory exists
        self._ensure_log_directory()

        # Load existing history if available
        if self.config.enable_logging:
            self._load_history()

        self.logger.info(
            f"PaperTradingManager initialized: "
            f"balances={self.config.initial_balances}, "
            f"fee_rate={self._trading_fee_rate:.4f}, "
            f"slippage={self.config.slippage_pct:.4f}%"
        )

    def _ensure_log_directory(self) -> None:
        """Create log directory if it doesn't exist."""
        log_path = Path(self.config.history_file)
        log_path.parent.mkdir(parents=True, exist_ok=True)

    def _load_history(self) -> None:
        """Load existing trade history from file."""
        history_path = Path(self.config.history_file)
        if not history_path.exists():
            return

        try:
            with open(history_path, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        entry = PaperTradeEntry.model_validate_json(line)
                        self._trade_history.append(entry)
                    except Exception as e:
                        self.logger.warning(f"Failed to parse history entry: {e}")

            self.logger.info(f"Loaded {len(self._trade_history)} paper trades from history")
        except Exception as e:
            self.logger.error(f"Failed to load trade history: {e}")

    @property
    def is_enabled(self) -> bool:
        """Check if paper trading is enabled."""
        return self.config.enabled

    @property
    def exchange(self) -> MockExchange:
        """Get the underlying mock exchange."""
        return self._exchange

    def enable(self) -> None:
        """Enable paper trading mode."""
        self.config.enabled = True
        self.logger.info("Paper trading enabled")

    def disable(self) -> None:
        """Disable paper trading mode."""
        self.config.enabled = False
        self.logger.info("Paper trading disabled")

    async def connect(self) -> bool:
        """Connect to paper trading exchange."""
        return await self._exchange.connect()

    async def disconnect(self) -> None:
        """Disconnect from paper trading exchange."""
        await self._exchange.disconnect()

    async def get_balance(self, asset: Optional[str] = None) -> List[Balance]:
        """Get virtual account balance.

        Args:
            asset: Optional asset filter

        Returns:
            List of balance objects
        """
        return await self._exchange.get_balance(asset)

    async def get_balance_by_asset(self, asset: str) -> Optional[Balance]:
        """Get balance for a specific asset.

        Args:
            asset: Asset symbol (e.g., "BTC", "USDT")

        Returns:
            Balance object or None if not found
        """
        balances = await self.get_balance(asset)
        for balance in balances:
            if balance.asset.upper() == asset.upper():
                return balance
        return None

    def get_balance_dict(self) -> Dict[str, float]:
        """Get balances as a simple dictionary.

        Returns:
            Dictionary of {asset: total_balance}
        """
        balances = asyncio.run(self.get_balance())
        return {b.asset: b.total for b in balances}

    async def place_market_order(
        self,
        symbol: str,
        side: OrderSide,
        amount: float,
    ) -> Order:
        """Place a market order in paper trading mode.

        Args:
            symbol: Trading symbol (e.g., "BTC/USDT")
            side: Buy or sell
            amount: Order amount

        Returns:
            Order object with execution details

        Raises:
            InsufficientFundsError: If virtual balance is insufficient
        """
        if not self.config.enabled:
            raise RuntimeError("Paper trading is disabled")

        # Record balances before trade
        balance_before = self.get_balance_dict()

        # Place order through mock exchange
        order = await self._exchange.place_order(
            symbol=symbol,
            side=side,
            amount=amount,
            order_type="market",
        )

        # Calculate fee
        trade_value = order.price * order.filled_amount
        fee = trade_value * self._trading_fee_rate

        # Record balances after trade
        balance_after = self.get_balance_dict()

        # Create trade entry
        entry = PaperTradeEntry(
            timestamp=datetime.now().isoformat(),
            symbol=symbol.upper(),
            side=side.value,
            order_type="market",
            amount=amount,
            price=order.price,
            filled_amount=order.filled_amount,
            fee=fee,
            slippage_pct=self.config.slippage_pct,
            balance_before=balance_before,
            balance_after=balance_after,
            order_id=order.order_id,
        )

        # Log and save
        self._record_trade(entry)

        self.logger.info(
            f"Paper market order: {side.value} {amount} {symbol} @ {order.price:.4f} "
            f"(fee: ${fee:.4f})"
        )

        return order

    async def place_limit_order(
        self,
        symbol: str,
        side: OrderSide,
        amount: float,
        price: float,
    ) -> Order:
        """Place a limit order in paper trading mode.

        For paper trading, limit orders are executed immediately if the
        price is within the current spread.

        Args:
            symbol: Trading symbol
            side: Buy or sell
            amount: Order amount
            price: Limit price

        Returns:
            Order object
        """
        if not self.config.enabled:
            raise RuntimeError("Paper trading is disabled")

        # Record balances before trade
        balance_before = self.get_balance_dict()

        # Get current ticker to check if limit price is achievable
        ticker = await self._exchange.get_ticker(symbol)

        # For paper trading, execute immediately if price is reasonable
        # Buy: execute if limit >= current ask
        # Sell: execute if limit <= current bid
        should_execute = False
        if side == OrderSide.BUY and price >= ticker.ask:
            should_execute = True
        elif side == OrderSide.SELL and price <= ticker.bid:
            should_execute = True

        if should_execute:
            order = await self._exchange.place_order(
                symbol=symbol,
                side=side,
                amount=amount,
                price=price,
                order_type="limit",
            )
        else:
            # Create pending order (not filled)
            order_id = str(uuid.uuid4())[:8]
            order = Order(
                order_id=order_id,
                symbol=symbol.upper(),
                side=side,
                order_type=OrderType.LIMIT,
                amount=amount,
                price=price,
                status=OrderStatus.OPEN,
                filled_amount=0.0,
            )

        # Calculate fee
        if order.filled_amount > 0:
            trade_value = order.price * order.filled_amount
            fee = trade_value * self._trading_fee_rate
        else:
            fee = 0.0

        # Record balances after trade
        balance_after = self.get_balance_dict()

        # Create trade entry (only if filled)
        if order.filled_amount > 0:
            entry = PaperTradeEntry(
                timestamp=datetime.now().isoformat(),
                symbol=symbol.upper(),
                side=side.value,
                order_type="limit",
                amount=amount,
                price=order.price,
                filled_amount=order.filled_amount,
                fee=fee,
                slippage_pct=0.0,  # Limit orders have no slippage
                balance_before=balance_before,
                balance_after=balance_after,
                order_id=order.order_id,
            )
            self._record_trade(entry)

        self.logger.info(
            f"Paper limit order: {side.value} {amount} {symbol} @ {price:.4f} "
            f"(filled: {order.filled_amount}, fee: ${fee:.4f})"
        )

        return order

    def _record_trade(self, entry: PaperTradeEntry) -> None:
        """Record a trade to history and file."""
        self._trade_history.append(entry)

        if self.config.enable_logging:
            self._write_trade_to_file(entry)

    def _write_trade_to_file(self, entry: PaperTradeEntry) -> None:
        """Write trade entry to history file."""
        try:
            with open(self.config.history_file, "a", encoding="utf-8") as f:
                f.write(entry.model_dump_json() + "\n")
        except Exception as e:
            self.logger.error(f"Failed to write trade history: {e}")

    async def cancel_order(self, order_id: str) -> bool:
        """Cancel an open order.

        Args:
            order_id: Order ID to cancel

        Returns:
            True if cancelled successfully
        """
        return await self._exchange.cancel_order(order_id)

    async def get_order(self, order_id: str) -> Optional[Order]:
        """Get order details.

        Args:
            order_id: Order ID to query

        Returns:
            Order object or None
        """
        return await self._exchange.get_order(order_id)

    async def get_open_orders(self, symbol: Optional[str] = None) -> List[Order]:
        """Get all open orders.

        Args:
            symbol: Optional symbol filter

        Returns:
            List of open orders
        """
        return await self._exchange.get_open_orders(symbol)

    async def get_positions(self, symbol: Optional[str] = None) -> List[Position]:
        """Get current positions.

        Args:
            symbol: Optional symbol filter

        Returns:
            List of positions
        """
        return await self._exchange.get_positions(symbol)

    async def get_ticker(self, symbol: str) -> Ticker:
        """Get current market ticker.

        Args:
            symbol: Trading symbol

        Returns:
            Ticker object with current prices
        """
        return await self._exchange.get_ticker(symbol)

    def update_ticker_price(self, symbol: str, price: float) -> None:
        """Manually update ticker price for testing.

        Args:
            symbol: Trading symbol
            price: New price
        """
        self._exchange.update_ticker(symbol, price)

    def set_virtual_balance(self, asset: str, amount: float) -> None:
        """Set virtual balance for an asset.

        Args:
            asset: Asset symbol
            amount: New balance amount
        """
        self._exchange.set_balance(asset, amount)
        self.logger.info(f"Virtual balance updated: {asset} = {amount}")

    def get_trade_history(
        self,
        symbol: Optional[str] = None,
        start_time: Optional[str] = None,
        end_time: Optional[str] = None,
    ) -> List[PaperTradeEntry]:
        """Get paper trade history with optional filtering.

        Args:
            symbol: Optional symbol filter
            start_time: Optional start time (ISO format)
            end_time: Optional end time (ISO format)

        Returns:
            List of trade entries
        """
        entries = self._trade_history.copy()

        if symbol:
            entries = [e for e in entries if e.symbol.upper() == symbol.upper()]

        if start_time:
            entries = [e for e in entries if e.timestamp >= start_time]

        if end_time:
            entries = [e for e in entries if e.timestamp <= end_time]

        return entries

    def get_performance_stats(self) -> Dict[str, Any]:
        """Get paper trading performance statistics.

        Returns:
            Dictionary with performance metrics
        """
        if not self._trade_history:
            return {
                "total_trades": 0,
                "total_volume": 0.0,
                "total_fees": 0.0,
                "win_rate": 0.0,
                "avg_trade_size": 0.0,
            }

        total_trades = len(self._trade_history)
        total_volume = sum(e.price * e.filled_amount for e in self._trade_history)
        total_fees = sum(e.fee for e in self._trade_history)

        # Calculate win rate (simplified - based on consecutive trades)
        wins = 0
        for i, entry in enumerate(self._trade_history):
            if entry.pnl > 0:
                wins += 1

        win_rate = wins / total_trades if total_trades > 0 else 0.0
        avg_trade_size = total_volume / total_trades if total_trades > 0 else 0.0

        # Current balances
        current_balances = self.get_balance_dict()
        initial_balances = self.config.initial_balances

        # Calculate total portfolio value change
        # This is simplified - real implementation would need current prices
        portfolio_change = sum(
            current_balances.get(asset, 0) - initial_balances.get(asset, 0)
            for asset in set(current_balances.keys()) | set(initial_balances.keys())
        )

        return {
            "total_trades": total_trades,
            "total_volume": round(total_volume, 4),
            "total_fees": round(total_fees, 4),
            "win_rate": round(win_rate * 100, 2),
            "avg_trade_size": round(avg_trade_size, 4),
            "current_balances": current_balances,
            "initial_balances": initial_balances,
            "portfolio_change": round(portfolio_change, 4),
        }

    def reset(self) -> None:
        """Reset paper trading to initial state.

        Clears all balances, positions, and history.
        """
        # Reset balances
        self._exchange._balances = {}
        for asset, amount in self.config.initial_balances.items():
            self._exchange._balances[asset.upper()] = Balance(
                asset=asset.upper(), free=amount, locked=0.0
            )

        # Clear orders and positions
        self._exchange._orders = {}
        self._exchange._positions = {}
        self._positions = {}
        self._trade_history = []

        self.logger.info("Paper trading reset to initial state")

    def export_history_to_csv(self, filepath: str) -> None:
        """Export trade history to CSV format.

        Args:
            filepath: Path to output CSV file
        """
        import csv

        path = Path(filepath)
        path.parent.mkdir(parents=True, exist_ok=True)

        with open(path, "w", newline="", encoding="utf-8") as f:
            writer = csv.writer(f)
            writer.writerow([
                "timestamp", "symbol", "side", "order_type", "amount",
                "price", "filled_amount", "fee", "slippage_pct", "order_id"
            ])

            for entry in self._trade_history:
                writer.writerow([
                    entry.timestamp,
                    entry.symbol,
                    entry.side,
                    entry.order_type,
                    entry.amount,
                    entry.price,
                    entry.filled_amount,
                    entry.fee,
                    entry.slippage_pct,
                    entry.order_id,
                ])

        self.logger.info(f"Trade history exported to {filepath}")

    def __repr__(self) -> str:
        """String representation."""
        return (
            f"PaperTradingManager("
            f"enabled={self.config.enabled}, "
            f"trades={len(self._trade_history)}, "
            f"fee_rate={self._trading_fee_rate:.4f})"
        )
