#!/usr/bin/env python3
"""
Trading Decision Prompts
========================

Prompt templates for the LLM trading decision agent.
Formats ML predictions and market context into structured prompts
that guide the LLM to make informed trading decisions.

Key Components:
- SYSTEM_PROMPT: Core identity and rules for the trading agent
- TradingDecisionPrompt: Builder class for constructing complete prompts

Author: ML-Specialist (NEXUS v4.0)
Version: 1.0.0
Created: 2026-01-25
"""

from typing import Dict, List, Optional, Any
from dataclasses import dataclass


# ============================================================
# SYSTEM PROMPT - Core Identity for Trading Agent
# ============================================================

SYSTEM_PROMPT = """You are an expert algorithmic trading agent specializing in ML-assisted decision making.

## Your Role
You analyze predictions from an ML ensemble system and decide whether to execute trades.
The ML system provides:
- Directional predictions (BULLISH/BEARISH/NEUTRAL)
- Magnitude predictions (expected price movement %)
- Confidence scores from multiple strategies
- Strategy weights learned from historical performance

## Risk Management Rules (MANDATORY)
1. Maximum 2% risk per trade ($20 on $1,000 account)
2. Maximum 2 simultaneous positions
3. Stop loss is ALWAYS required
4. Minimum risk:reward ratio of 1.5:1
5. Maximum drawdown allowed: 15%
6. Never trade against strong consensus (3+ strategies agreeing)

## Decision Framework
For each signal, evaluate:
1. ML Confidence: Is ensemble confidence above threshold (60%)?
2. Strategy Alignment: Do multiple strategies agree on direction?
3. Risk/Reward: Does the predicted magnitude justify the stop distance?
4. Market Context: Is the current volatility suitable for trading?
5. Account Status: Can we afford this risk given current drawdown?

## Output Format
You MUST respond with a JSON object containing your decision.
Do not include any text outside the JSON block.

```json
{
    "decision": "TRADE" | "NO_TRADE" | "WAIT",
    "direction": "LONG" | "SHORT" | null,
    "entry": <price or "MARKET">,
    "stop_loss": <price>,
    "take_profit": [<tp1>, <tp2>],
    "position_size": <percentage of account, e.g., 1.0>,
    "reasoning": "<brief explanation of your decision>"
}
```

## Decision Types
- TRADE: Execute the trade with specified parameters
- NO_TRADE: Do not trade - conditions are unfavorable
- WAIT: Conditions are promising but need confirmation

Be concise, precise, and always prioritize capital preservation."""


# ============================================================
# EXPECTED RESPONSE FORMAT
# ============================================================

EXPECTED_RESPONSE_FORMAT = """{
    "decision": "TRADE" | "NO_TRADE" | "WAIT",
    "direction": "LONG" | "SHORT" | null,
    "entry": <price or "MARKET">,
    "stop_loss": <price>,
    "take_profit": [<tp1>, <tp2>],
    "position_size": <percentage>,
    "reasoning": "<explanation>"
}"""


# ============================================================
# PROMPT TEMPLATES
# ============================================================

PREDICTION_SUMMARY_TEMPLATE = """## ML Ensemble Prediction for {symbol}
**Timestamp:** {timestamp}
**Timeframe:** {timeframe}

### Ensemble Output
- **Direction:** {direction} (net score: {net_score:+.2f})
- **Predicted Move:** {magnitude}
- **Ensemble Confidence:** {confidence:.1%}
- **Strategies Aligned:** {aligned_count}/{total_strategies}
"""

STRATEGY_TABLE_TEMPLATE = """### Individual Strategy Signals
| Strategy | Direction | Confidence | Weight |
|----------|-----------|------------|--------|
{rows}
"""

MARKET_CONTEXT_TEMPLATE = """### Market Context
- **Current Price:** {current_price}
- **ATR (14):** {atr} ({atr_pct:.2%} of price)
- **Volatility Regime:** {volatility_regime}
- **Market Phase:** {market_phase}
- **Session:** {trading_session}
"""

RISK_PARAMETERS_TEMPLATE = """### Account & Risk Parameters
- **Account Balance:** ${balance:.2f}
- **Available Margin:** ${available_margin:.2f}
- **Current Drawdown:** {drawdown:.1%}
- **Open Positions:** {open_positions}
- **Max Risk Per Trade:** {max_risk_pct:.1%} (${max_risk_usd:.2f})
- **Today's P&L:** ${daily_pnl:+.2f} ({daily_pnl_pct:+.1%})
"""

DECISION_TASK_TEMPLATE = """### Your Task
Based on the ML predictions and market context above, decide whether to:
1. **TRADE** - Enter a position with your recommended parameters
2. **NO_TRADE** - Skip this signal (explain why)
3. **WAIT** - Signal is promising but needs confirmation

If you decide to TRADE, provide:
- Entry price (or "MARKET" for market order)
- Stop loss level
- Take profit target(s)
- Position size as % of account to risk

Remember: Capital preservation is paramount. When in doubt, NO_TRADE.
"""


# ============================================================
# TRADING DECISION PROMPT BUILDER
# ============================================================

@dataclass
class MetamodelOutput:
    """Output from the metamodel"""
    direction: str  # "BULLISH", "BEARISH", "NEUTRAL"
    magnitude: float  # Expected move in percentage
    confidence: float  # 0-1 confidence score
    strategy_weights: Dict[str, float]  # Strategy name -> weight
    delta_high: float  # Predicted high delta
    delta_low: float  # Predicted low delta


@dataclass
class StrategyOutput:
    """Output from an individual strategy"""
    name: str
    direction: str  # "BULLISH", "BEARISH", "NEUTRAL"
    confidence: float  # 0-1
    weight: float  # Current weight from metamodel
    details: Optional[Dict[str, Any]] = None


@dataclass
class MarketContext:
    """Current market context"""
    current_price: float
    atr: float
    volatility_regime: str  # "low", "medium", "high"
    market_phase: str  # "accumulation", "distribution", "manipulation", "trend"
    trading_session: str  # "asian", "london", "new_york", "overlap"
    timestamp: str
    timeframe: str


@dataclass
class AccountInfo:
    """Account and risk information"""
    balance: float
    available_margin: float
    current_drawdown: float
    open_positions: int
    max_risk_per_trade: float  # Percentage
    daily_pnl: float
    daily_pnl_pct: float


class TradingDecisionPrompt:
    """
    Builder class for constructing trading decision prompts.

    Formats ML predictions and market context into a structured prompt
    that guides the LLM to make informed trading decisions.

    Usage:
        prompt_builder = TradingDecisionPrompt()
        full_prompt = prompt_builder.build_full_prompt(
            symbol="XAUUSD",
            metamodel_output=metamodel_result,
            strategy_outputs=strategy_results,
            market_context=context,
            account_info=account
        )
    """

    def __init__(self, system_prompt: str = None):
        """
        Initialize the prompt builder.

        Args:
            system_prompt: Custom system prompt (uses default if None)
        """
        self.system_prompt = system_prompt or SYSTEM_PROMPT

    def format_prediction_summary(
        self,
        symbol: str,
        metamodel_output: MetamodelOutput,
        strategy_outputs: List[StrategyOutput],
        timestamp: str = "now",
        timeframe: str = "5m"
    ) -> str:
        """
        Format the ML prediction summary section.

        Args:
            symbol: Trading symbol (e.g., "XAUUSD")
            metamodel_output: Output from the metamodel
            strategy_outputs: List of individual strategy outputs
            timestamp: Current timestamp
            timeframe: Analysis timeframe

        Returns:
            Formatted prediction summary string
        """
        # Calculate net score based on direction
        direction_scores = {
            "BULLISH": 1.0,
            "BEARISH": -1.0,
            "NEUTRAL": 0.0
        }
        net_score = direction_scores.get(metamodel_output.direction, 0.0) * metamodel_output.confidence

        # Count aligned strategies
        aligned_count = sum(
            1 for s in strategy_outputs
            if s.direction == metamodel_output.direction and s.confidence > 0.5
        )

        # Format magnitude
        if metamodel_output.direction == "BULLISH":
            magnitude = f"+{metamodel_output.magnitude:.2f}%"
        elif metamodel_output.direction == "BEARISH":
            magnitude = f"-{metamodel_output.magnitude:.2f}%"
        else:
            magnitude = f"{metamodel_output.magnitude:.2f}% (range)"

        return PREDICTION_SUMMARY_TEMPLATE.format(
            symbol=symbol,
            timestamp=timestamp,
            timeframe=timeframe,
            direction=metamodel_output.direction,
            net_score=net_score,
            magnitude=magnitude,
            confidence=metamodel_output.confidence,
            aligned_count=aligned_count,
            total_strategies=len(strategy_outputs)
        )

    def format_strategy_table(
        self,
        strategy_outputs: List[StrategyOutput],
        weights: Optional[Dict[str, float]] = None
    ) -> str:
        """
        Format the strategy signals as a markdown table.

        Args:
            strategy_outputs: List of individual strategy outputs
            weights: Optional weight override (uses strategy.weight if None)

        Returns:
            Formatted markdown table string
        """
        rows = []
        for strategy in strategy_outputs:
            weight = weights.get(strategy.name, strategy.weight) if weights else strategy.weight

            # Format direction with emoji indicator
            direction_display = {
                "BULLISH": "BULLISH",
                "BEARISH": "BEARISH",
                "NEUTRAL": "NEUTRAL"
            }.get(strategy.direction, strategy.direction)

            row = f"| {strategy.name} | {direction_display} | {strategy.confidence:.1%} | {weight:.1%} |"
            rows.append(row)

        return STRATEGY_TABLE_TEMPLATE.format(rows="\n".join(rows))

    def format_market_context(self, context: MarketContext) -> str:
        """
        Format the market context section.

        Args:
            context: MarketContext dataclass with current market info

        Returns:
            Formatted market context string
        """
        atr_pct = context.atr / context.current_price if context.current_price > 0 else 0

        return MARKET_CONTEXT_TEMPLATE.format(
            current_price=f"{context.current_price:.4f}" if context.current_price < 100 else f"{context.current_price:.2f}",
            atr=f"{context.atr:.4f}" if context.atr < 1 else f"{context.atr:.2f}",
            atr_pct=atr_pct,
            volatility_regime=context.volatility_regime.upper(),
            market_phase=context.market_phase.title(),
            trading_session=context.trading_session.replace("_", " ").title()
        )

    def format_risk_parameters(self, account_info: AccountInfo) -> str:
        """
        Format the risk parameters section.

        Args:
            account_info: AccountInfo dataclass with account details

        Returns:
            Formatted risk parameters string
        """
        max_risk_usd = account_info.balance * (account_info.max_risk_per_trade / 100)

        return RISK_PARAMETERS_TEMPLATE.format(
            balance=account_info.balance,
            available_margin=account_info.available_margin,
            drawdown=account_info.current_drawdown,
            open_positions=account_info.open_positions,
            max_risk_pct=account_info.max_risk_per_trade,
            max_risk_usd=max_risk_usd,
            daily_pnl=account_info.daily_pnl,
            daily_pnl_pct=account_info.daily_pnl_pct
        )

    def build_full_prompt(
        self,
        symbol: str,
        metamodel_output: MetamodelOutput,
        strategy_outputs: List[StrategyOutput],
        market_context: MarketContext,
        account_info: AccountInfo
    ) -> str:
        """
        Build the complete prompt for the LLM.

        Combines all sections into a comprehensive prompt that includes:
        - ML prediction summary
        - Strategy table
        - Market context
        - Risk parameters
        - Decision task

        Args:
            symbol: Trading symbol
            metamodel_output: Output from the metamodel
            strategy_outputs: List of individual strategy outputs
            market_context: Current market context
            account_info: Account and risk information

        Returns:
            Complete formatted prompt string
        """
        sections = [
            self.format_prediction_summary(
                symbol=symbol,
                metamodel_output=metamodel_output,
                strategy_outputs=strategy_outputs,
                timestamp=market_context.timestamp,
                timeframe=market_context.timeframe
            ),
            self.format_strategy_table(
                strategy_outputs=strategy_outputs,
                weights=metamodel_output.strategy_weights
            ),
            self.format_market_context(market_context),
            self.format_risk_parameters(account_info),
            DECISION_TASK_TEMPLATE
        ]

        return "\n".join(sections)

    def get_system_prompt(self) -> str:
        """Get the system prompt for the LLM."""
        return self.system_prompt

    @staticmethod
    def create_from_dict(
        symbol: str,
        metamodel_dict: Dict[str, Any],
        strategies_dict: List[Dict[str, Any]],
        context_dict: Dict[str, Any],
        account_dict: Dict[str, Any]
    ) -> str:
        """
        Convenience method to create a prompt from dictionaries.

        Useful when data comes from JSON/API responses.

        Args:
            symbol: Trading symbol
            metamodel_dict: Metamodel output as dict
            strategies_dict: Strategy outputs as list of dicts
            context_dict: Market context as dict
            account_dict: Account info as dict

        Returns:
            Complete formatted prompt string
        """
        builder = TradingDecisionPrompt()

        metamodel = MetamodelOutput(
            direction=metamodel_dict.get("direction", "NEUTRAL"),
            magnitude=metamodel_dict.get("magnitude", 0.0),
            confidence=metamodel_dict.get("confidence", 0.5),
            strategy_weights=metamodel_dict.get("strategy_weights", {}),
            delta_high=metamodel_dict.get("delta_high", 0.0),
            delta_low=metamodel_dict.get("delta_low", 0.0)
        )

        strategies = [
            StrategyOutput(
                name=s.get("name", "Unknown"),
                direction=s.get("direction", "NEUTRAL"),
                confidence=s.get("confidence", 0.5),
                weight=s.get("weight", 0.2),
                details=s.get("details")
            )
            for s in strategies_dict
        ]

        context = MarketContext(
            current_price=context_dict.get("current_price", 0.0),
            atr=context_dict.get("atr", 0.0),
            volatility_regime=context_dict.get("volatility_regime", "medium"),
            market_phase=context_dict.get("market_phase", "unknown"),
            trading_session=context_dict.get("trading_session", "unknown"),
            timestamp=context_dict.get("timestamp", "now"),
            timeframe=context_dict.get("timeframe", "5m")
        )

        account = AccountInfo(
            balance=account_dict.get("balance", 1000.0),
            available_margin=account_dict.get("available_margin", 1000.0),
            current_drawdown=account_dict.get("current_drawdown", 0.0),
            open_positions=account_dict.get("open_positions", 0),
            max_risk_per_trade=account_dict.get("max_risk_per_trade", 2.0),
            daily_pnl=account_dict.get("daily_pnl", 0.0),
            daily_pnl_pct=account_dict.get("daily_pnl_pct", 0.0)
        )

        return builder.build_full_prompt(
            symbol=symbol,
            metamodel_output=metamodel,
            strategy_outputs=strategies,
            market_context=context,
            account_info=account
        )


# ============================================================
# EXPORTS
# ============================================================

__all__ = [
    'SYSTEM_PROMPT',
    'EXPECTED_RESPONSE_FORMAT',
    'TradingDecisionPrompt',
    'MetamodelOutput',
    'StrategyOutput',
    'MarketContext',
    'AccountInfo',
]
