"""Economic tracker for trading agents.

This module provides the TradingEconomicTracker class for tracking agent
financial status, costs, and survival state in a trading environment.
"""

from datetime import datetime
from enum import Enum
from pathlib import Path

from pydantic import BaseModel, Field


class SurvivalStatus(str, Enum):
    """Agent survival status levels."""

    THRIVING = "🚀 thriving"
    STABLE = "💪 stable"
    STRUGGLING = "⚠️ struggling"
    CRITICAL = "🔴 critical"
    BANKRUPT = "💀 bankrupt"


class BalanceHistoryEntry(BaseModel):
    """Single balance history entry."""

    timestamp: str = Field(..., description="ISO format timestamp")
    balance: float = Field(..., ge=0, description="Balance after change")
    change: float = Field(..., description="Balance change amount")
    reason: str = Field(..., description="Reason for balance change")


class TradeCostResult(BaseModel):
    """Result of trade cost calculation."""

    fee: float = Field(..., ge=0, description="Trading fee paid")
    pnl: float = Field(..., description="Profit/loss from trade")
    balance: float = Field(..., ge=0, description="Current balance")
    status: SurvivalStatus = Field(..., description="Current survival status")


class EconomicTrackerState(BaseModel):
    """Complete state of the economic tracker for persistence."""

    agent_id: str
    initial_capital: float
    balance: float
    token_costs: float
    trade_costs: float
    realized_pnl: float
    thresholds: dict[str, float]
    token_cost_per_1m_input: float
    token_cost_per_1m_output: float
    trade_fee_rate: float
    data_cost_per_call: float
    balance_history: list[BalanceHistoryEntry]


class TradingEconomicTracker:
    """Track agent economic status and survival state.

    Each agent must pay for its own decisions and trades. The tracker
    monitors balance, costs, and determines survival status based on
    configurable thresholds relative to initial capital.

    Args:
        agent_id: Unique identifier for the agent
        initial_capital: Starting balance (default: $10,000)
        token_cost_per_1m_input: Cost per 1M input tokens (default: $2.5)
        token_cost_per_1m_output: Cost per 1M output tokens (default: $10.0)
        trade_fee_rate: Trading fee as decimal (default: 0.001 = 0.1%)
        data_cost_per_call: Cost per market data API call (default: $0.01)
    """

    # Default survival thresholds as multipliers of initial capital
    DEFAULT_THRESHOLDS = {
        "thriving": 1.5,  # 50% profit
        "stable": 1.1,  # 10% profit
        "struggling": 0.8,  # 20% loss
        "bankrupt": 0.3,  # 70% loss
    }

    def __init__(
        self,
        agent_id: str,
        initial_capital: float = 10000.0,
        token_cost_per_1m_input: float = 2.5,
        token_cost_per_1m_output: float = 10.0,
        trade_fee_rate: float = 0.001,
        data_cost_per_call: float = 0.01,
    ) -> None:
        self.agent_id: str = agent_id
        self.initial_capital: float = initial_capital
        self.balance: float = initial_capital
        self.token_costs: float = 0.0
        self.trade_costs: float = 0.0
        self.realized_pnl: float = 0.0

        # Cost parameters
        self.token_cost_per_1m_input: float = token_cost_per_1m_input
        self.token_cost_per_1m_output: float = token_cost_per_1m_output
        self.trade_fee_rate: float = trade_fee_rate
        self.data_cost_per_call: float = data_cost_per_call

        # Calculate absolute thresholds
        self.thresholds: dict[str, float] = {
            key: initial_capital * multiplier
            for key, multiplier in self.DEFAULT_THRESHOLDS.items()
        }

        # Balance history tracking
        self._balance_history: list[BalanceHistoryEntry] = [
            BalanceHistoryEntry(
                timestamp=datetime.now().isoformat(),
                balance=initial_capital,
                change=0.0,
                reason="Initial capital",
            )
        ]

    @property
    def total_costs(self) -> float:
        """Return total accumulated costs (token + trade fees)."""
        return self.token_costs + self.trade_costs

    @property
    def net_profit(self) -> float:
        """Return net profit (realized PnL minus total costs)."""
        return self.realized_pnl - self.total_costs

    def calculate_decision_cost(
        self,
        tokens_input: int,
        tokens_output: int,
        market_data_calls: int = 0,
    ) -> float:
        """Calculate and deduct the cost of a decision.

        Costs include:
        - LLM token costs (input and output)
        - Market data API call costs

        Args:
            tokens_input: Number of input tokens to LLM
            tokens_output: Number of output tokens from LLM
            market_data_calls: Number of market data API calls

        Returns:
            Total cost of the decision (rounded to 4 decimal places)
        """
        # LLM token costs (per million tokens)
        llm_cost = (
            tokens_input / 1e6 * self.token_cost_per_1m_input
            + tokens_output / 1e6 * self.token_cost_per_1m_output
        )

        # Market data costs
        data_cost = market_data_calls * self.data_cost_per_call

        total_cost = round(llm_cost + data_cost, 4)

        # Update state
        self.token_costs += total_cost
        self._update_balance(-total_cost, f"Decision cost: {tokens_input}in/{tokens_output}out/{market_data_calls}calls")

        return total_cost

    def calculate_trade_cost(
        self,
        trade_value: float,
        is_win: bool,
        win_amount: float = 0.0,
        loss_amount: float = 0.0,
    ) -> TradeCostResult:
        """Calculate and apply trade costs and PnL.

        Args:
            trade_value: Total value of the trade
            is_win: Whether the trade was profitable
            win_amount: Profit amount (if win)
            loss_amount: Loss amount (if loss)

        Returns:
            TradeCostResult with fee, PnL, balance, and status
        """
        # Calculate trading fee
        fee = round(trade_value * self.trade_fee_rate, 4)
        self.trade_costs += fee

        # Calculate PnL (win - loss - fee)
        gross_pnl = win_amount - loss_amount
        pnl = round(gross_pnl - fee, 4)
        self.realized_pnl += gross_pnl

        # Update balance: deduct fee first, then apply PnL
        net_change = pnl  # pnl already includes -fee
        self._update_balance(net_change, f"Trade: {'win' if is_win else 'loss'} ${abs(gross_pnl):.2f}")

        return TradeCostResult(
            fee=fee,
            pnl=pnl,
            balance=self.balance,
            status=self.get_survival_status(),
        )

    def deduct(self, amount: float, reason: str) -> bool:
        """Deduct a fixed amount from balance (e.g. factor purchase, course enrollment).

        Args:
            amount: Positive amount to deduct
            reason: Description for history

        Returns:
            True if deduction was applied, False if insufficient balance
        """
        if amount <= 0:
            return True
        if self.balance < amount:
            return False
        self._update_balance(-amount, reason)
        return True

    def _update_balance(self, change: float, reason: str) -> None:
        """Update balance and record in history.

        Args:
            change: Amount to add (positive) or subtract (negative)
            reason: Description of the balance change
        """
        self.balance = round(max(0.0, self.balance + change), 4)

        entry = BalanceHistoryEntry(
            timestamp=datetime.now().isoformat(),
            balance=self.balance,
            change=round(change, 4),
            reason=reason,
        )
        self._balance_history.append(entry)

    def get_survival_status(self) -> SurvivalStatus:
        """Determine current survival status based on balance.

        Status levels (from highest to lowest balance):
        - THRIVING: Balance >= 150% of initial capital
        - STABLE: Balance >= 110% of initial capital
        - STRUGGLING: Balance >= 80% of initial capital
        - CRITICAL: Balance >= 30% of initial capital
        - BANKRUPT: Balance < 30% of initial capital

        Returns:
            Current SurvivalStatus enum value
        """
        if self.balance >= self.thresholds["thriving"]:
            return SurvivalStatus.THRIVING
        elif self.balance >= self.thresholds["stable"]:
            return SurvivalStatus.STABLE
        elif self.balance >= self.thresholds["struggling"]:
            return SurvivalStatus.STRUGGLING
        elif self.balance >= self.thresholds["bankrupt"]:
            return SurvivalStatus.CRITICAL
        else:
            return SurvivalStatus.BANKRUPT

    def get_balance_history(self) -> list[BalanceHistoryEntry]:
        """Return complete balance history.

        Returns:
            List of BalanceHistoryEntry from initial capital to present
        """
        return self._balance_history.copy()

    def get_state(self) -> EconomicTrackerState:
        """Get complete state for persistence.

        Returns:
            EconomicTrackerState with all tracker data
        """
        return EconomicTrackerState(
            agent_id=self.agent_id,
            initial_capital=self.initial_capital,
            balance=self.balance,
            token_costs=self.token_costs,
            trade_costs=self.trade_costs,
            realized_pnl=self.realized_pnl,
            thresholds=self.thresholds.copy(),
            token_cost_per_1m_input=self.token_cost_per_1m_input,
            token_cost_per_1m_output=self.token_cost_per_1m_output,
            trade_fee_rate=self.trade_fee_rate,
            data_cost_per_call=self.data_cost_per_call,
            balance_history=self._balance_history.copy(),
        )

    def save_to_file(self, filepath: str | Path) -> None:
        """Save tracker state to JSONL file.

        Each line is a JSON object representing a state snapshot.

        Args:
            filepath: Path to save the state file
        """
        path = Path(filepath)
        path.parent.mkdir(parents=True, exist_ok=True)

        state = self.get_state()

        with open(path, "a", encoding="utf-8") as f:
            f.write(state.model_dump_json() + "\n")

    @classmethod
    def load_from_file(cls, filepath: str | Path) -> "TradingEconomicTracker":
        """Load tracker from latest state in JSONL file.

        Args:
            filepath: Path to the state file

        Returns:
            TradingEconomicTracker restored from file

        Raises:
            FileNotFoundError: If file doesn't exist
            ValueError: If file is empty or contains invalid data
        """
        path = Path(filepath)

        if not path.exists():
            raise FileNotFoundError(f"State file not found: {filepath}")

        lines = path.read_text(encoding="utf-8").strip().split("\n")
        if not lines or lines == [""]:
            raise ValueError(f"State file is empty: {filepath}")

        # Load latest state (last line)
        latest_line = lines[-1]
        state_data = EconomicTrackerState.model_validate_json(latest_line)

        # Create new tracker with loaded parameters
        tracker = cls(
            agent_id=state_data.agent_id,
            initial_capital=state_data.initial_capital,
            token_cost_per_1m_input=state_data.token_cost_per_1m_input,
            token_cost_per_1m_output=state_data.token_cost_per_1m_output,
            trade_fee_rate=state_data.trade_fee_rate,
            data_cost_per_call=state_data.data_cost_per_call,
        )

        # Restore state
        tracker.balance = state_data.balance
        tracker.token_costs = state_data.token_costs
        tracker.trade_costs = state_data.trade_costs
        tracker.realized_pnl = state_data.realized_pnl
        tracker.thresholds = state_data.thresholds
        tracker._balance_history = state_data.balance_history

        return tracker

    def __repr__(self) -> str:
        """Return string representation of tracker state."""
        return (
            f"TradingEconomicTracker("
            f"agent_id='{self.agent_id}', "
            f"balance=${self.balance:.2f}, "
            f"status={self.get_survival_status().value}, "
            f"total_costs=${self.total_costs:.2f}, "
            f"pnl=${self.realized_pnl:.2f})"
        )
