#!/usr/bin/env python3
"""
Decision Parser
===============

Parses LLM responses into structured trading decisions.
Extracts decision components (action, entry, stop loss, take profit)
from the LLM's JSON response and validates them.

Key Components:
- TradingDecision: Dataclass representing a parsed decision
- DecisionParser: Parser with extraction and validation logic

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

import re
import json
from typing import Dict, List, Optional, Any, Union, Tuple
from dataclasses import dataclass, field, asdict
from loguru import logger


# ============================================================
# DATA CLASSES
# ============================================================

@dataclass
class TradingDecision:
    """
    Structured trading decision extracted from LLM response.

    Represents the complete output of the LLM trading agent,
    including the action to take and all trade parameters.
    """
    # Core decision
    action: str  # "TRADE", "NO_TRADE", "WAIT"
    direction: Optional[str] = None  # "LONG", "SHORT", None

    # Trade parameters
    entry: Optional[Union[float, str]] = None  # Price or "MARKET"
    stop_loss: Optional[float] = None
    take_profit: List[float] = field(default_factory=list)
    position_size: float = 1.0  # Percentage of account to risk

    # Reasoning
    reasoning: str = ""

    # Raw response for debugging
    raw_response: str = ""

    # Validation status
    is_valid: bool = True
    validation_errors: List[str] = field(default_factory=list)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to serializable dictionary"""
        return asdict(self)

    def get_primary_take_profit(self) -> Optional[float]:
        """Get the first (primary) take profit level"""
        return self.take_profit[0] if self.take_profit else None

    def get_risk_reward_ratio(self) -> Optional[float]:
        """
        Calculate the risk/reward ratio if entry, SL, and TP are available.

        Returns:
            Risk/reward ratio or None if cannot be calculated
        """
        if not all([self.entry, self.stop_loss, self.take_profit]):
            return None

        entry = self.entry if isinstance(self.entry, (int, float)) else None
        if entry is None:
            return None

        tp = self.take_profit[0]
        sl = self.stop_loss

        risk = abs(entry - sl)
        reward = abs(tp - entry)

        if risk <= 0:
            return None

        return reward / risk


# ============================================================
# DECISION PARSER CLASS
# ============================================================

class DecisionParser:
    """
    Parser for LLM trading decision responses.

    Extracts structured trading decisions from LLM responses,
    handling various response formats and performing validation.

    Usage:
        parser = DecisionParser()

        # Parse LLM response
        decision = parser.parse_response(llm_response_text)

        # Check if valid
        if decision.is_valid:
            execute_trade(decision)
        else:
            logger.warning(f"Invalid decision: {decision.validation_errors}")
    """

    # Valid values for enum fields
    VALID_ACTIONS = {"TRADE", "NO_TRADE", "WAIT"}
    VALID_DIRECTIONS = {"LONG", "SHORT", None}

    # Validation rules
    MIN_POSITION_SIZE = 0.1  # 0.1%
    MAX_POSITION_SIZE = 5.0  # 5%
    MIN_RISK_REWARD = 1.0  # Minimum acceptable R:R

    def __init__(
        self,
        strict_mode: bool = False,
        min_risk_reward: float = 1.0,
        max_position_size: float = 5.0
    ):
        """
        Initialize the decision parser.

        Args:
            strict_mode: If True, reject decisions with any validation warnings
            min_risk_reward: Minimum acceptable risk/reward ratio
            max_position_size: Maximum position size as percentage
        """
        self.strict_mode = strict_mode
        self.min_risk_reward = min_risk_reward
        self.max_position_size = max_position_size

    def parse_response(self, llm_response: str) -> TradingDecision:
        """
        Parse an LLM response into a TradingDecision.

        Handles various response formats:
        - Pure JSON
        - JSON wrapped in markdown code blocks
        - JSON with surrounding text

        Args:
            llm_response: Raw text response from the LLM

        Returns:
            TradingDecision with extracted values (check is_valid)
        """
        # Store raw response
        raw = llm_response

        # Try to extract JSON from the response
        json_str = self._extract_json(llm_response)

        if not json_str:
            return TradingDecision(
                action="NO_TRADE",
                reasoning="Failed to parse LLM response - no valid JSON found",
                raw_response=raw,
                is_valid=False,
                validation_errors=["No valid JSON found in response"]
            )

        try:
            data = json.loads(json_str)
        except json.JSONDecodeError as e:
            return TradingDecision(
                action="NO_TRADE",
                reasoning=f"JSON parsing error: {e}",
                raw_response=raw,
                is_valid=False,
                validation_errors=[f"JSON decode error: {str(e)}"]
            )

        # Extract fields
        decision = self._extract_decision(data)
        entry = self._extract_entry(data)
        stop_loss = self._extract_stop_loss(data)
        take_profit = self._extract_take_profit(data)
        position_size = self._extract_position_size(data)
        reasoning = self._extract_reasoning(data)
        direction = self._extract_direction(data)

        # Build TradingDecision
        trading_decision = TradingDecision(
            action=decision,
            direction=direction,
            entry=entry,
            stop_loss=stop_loss,
            take_profit=take_profit,
            position_size=position_size,
            reasoning=reasoning,
            raw_response=raw
        )

        # Validate
        self._validate_decision(trading_decision)

        return trading_decision

    def _extract_json(self, text: str) -> Optional[str]:
        """
        Extract JSON from text, handling various formats.

        Tries multiple extraction strategies:
        1. Pure JSON (entire text is JSON)
        2. JSON in markdown code block
        3. JSON object anywhere in text

        Args:
            text: Text potentially containing JSON

        Returns:
            Extracted JSON string or None
        """
        text = text.strip()

        # Strategy 1: Pure JSON
        if text.startswith("{") and text.endswith("}"):
            return text

        # Strategy 2: Markdown code block
        # Match ```json ... ``` or ``` ... ```
        code_block_pattern = r"```(?:json)?\s*(\{[\s\S]*?\})\s*```"
        match = re.search(code_block_pattern, text)
        if match:
            return match.group(1).strip()

        # Strategy 3: Find JSON object anywhere
        # Look for { ... } pattern
        brace_pattern = r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}"
        match = re.search(brace_pattern, text)
        if match:
            return match.group(0)

        return None

    def _extract_decision(self, data: Dict[str, Any]) -> str:
        """Extract and normalize the decision/action field"""
        # Try multiple field names
        for field_name in ["decision", "action", "trade_action"]:
            if field_name in data:
                value = str(data[field_name]).upper().strip()
                # Normalize variations
                if value in ("BUY", "SELL", "ENTER", "EXECUTE"):
                    return "TRADE"
                if value in ("SKIP", "PASS", "REJECT", "NO"):
                    return "NO_TRADE"
                if value in self.VALID_ACTIONS:
                    return value

        return "NO_TRADE"

    def _extract_direction(self, data: Dict[str, Any]) -> Optional[str]:
        """Extract and normalize the direction field"""
        for field_name in ["direction", "side", "trade_direction", "type"]:
            if field_name in data and data[field_name] is not None:
                value = str(data[field_name]).upper().strip()
                # Normalize variations
                if value in ("BUY", "BULLISH", "UP"):
                    return "LONG"
                if value in ("SELL", "BEARISH", "DOWN"):
                    return "SHORT"
                if value in ("LONG", "SHORT"):
                    return value

        return None

    def _extract_entry(self, data: Dict[str, Any]) -> Optional[Union[float, str]]:
        """Extract the entry price or "MARKET" """
        for field_name in ["entry", "entry_price", "price", "entry_level"]:
            if field_name in data:
                value = data[field_name]

                if value is None:
                    continue

                # Check for market order
                if isinstance(value, str):
                    if value.upper() in ("MARKET", "MKT", "NOW"):
                        return "MARKET"
                    try:
                        return float(value)
                    except ValueError:
                        continue

                if isinstance(value, (int, float)):
                    return float(value)

        return None

    def _extract_stop_loss(self, data: Dict[str, Any]) -> Optional[float]:
        """Extract the stop loss price"""
        for field_name in ["stop_loss", "sl", "stop", "stoploss"]:
            if field_name in data and data[field_name] is not None:
                try:
                    return float(data[field_name])
                except (ValueError, TypeError):
                    continue

        return None

    def _extract_take_profit(self, data: Dict[str, Any]) -> List[float]:
        """Extract take profit level(s)"""
        take_profits = []

        # Try array field first
        for field_name in ["take_profit", "tp", "take_profits", "targets"]:
            if field_name in data:
                value = data[field_name]

                if value is None:
                    continue

                # Handle list
                if isinstance(value, list):
                    for tp in value:
                        try:
                            take_profits.append(float(tp))
                        except (ValueError, TypeError):
                            continue

                # Handle single value
                elif isinstance(value, (int, float)):
                    take_profits.append(float(value))

                elif isinstance(value, str):
                    try:
                        take_profits.append(float(value))
                    except ValueError:
                        pass

        # Also check individual TP fields
        for i in range(1, 4):
            for field_name in [f"take_profit_{i}", f"tp{i}", f"tp_{i}"]:
                if field_name in data and data[field_name] is not None:
                    try:
                        take_profits.append(float(data[field_name]))
                    except (ValueError, TypeError):
                        continue

        return take_profits

    def _extract_position_size(self, data: Dict[str, Any]) -> float:
        """Extract position size as percentage"""
        for field_name in ["position_size", "size", "risk", "risk_percent", "risk_pct"]:
            if field_name in data and data[field_name] is not None:
                try:
                    value = float(data[field_name])
                    # Clamp to valid range
                    return max(self.MIN_POSITION_SIZE, min(self.max_position_size, value))
                except (ValueError, TypeError):
                    continue

        return 1.0  # Default 1%

    def _extract_reasoning(self, data: Dict[str, Any]) -> str:
        """Extract the reasoning/explanation"""
        for field_name in ["reasoning", "reason", "explanation", "rationale", "analysis"]:
            if field_name in data and data[field_name] is not None:
                return str(data[field_name])

        return ""

    def _validate_decision(self, decision: TradingDecision) -> None:
        """
        Validate the parsed decision and populate validation errors.

        Performs sanity checks on the decision:
        - Action is valid
        - TRADE decisions have required fields
        - Risk/reward is acceptable
        - Position size is within limits

        Args:
            decision: TradingDecision to validate (modified in place)
        """
        errors = []

        # Check action is valid
        if decision.action not in self.VALID_ACTIONS:
            errors.append(f"Invalid action: {decision.action}")
            decision.action = "NO_TRADE"

        # For TRADE decisions, validate required fields
        if decision.action == "TRADE":
            # Direction is required
            if decision.direction is None:
                errors.append("TRADE decision missing direction")

            # Stop loss is required
            if decision.stop_loss is None:
                errors.append("TRADE decision missing stop_loss")

            # At least one take profit is recommended
            if not decision.take_profit:
                errors.append("TRADE decision missing take_profit (warning)")

            # Validate risk/reward if we have the data
            rr = decision.get_risk_reward_ratio()
            if rr is not None and rr < self.min_risk_reward:
                errors.append(f"Risk/reward ratio {rr:.2f} below minimum {self.min_risk_reward}")

            # Validate entry/SL/TP relationship for LONG
            if decision.direction == "LONG" and decision.entry and decision.stop_loss:
                entry = decision.entry if isinstance(decision.entry, (int, float)) else None
                if entry and decision.stop_loss >= entry:
                    errors.append("LONG: stop_loss should be below entry")

            # Validate entry/SL/TP relationship for SHORT
            if decision.direction == "SHORT" and decision.entry and decision.stop_loss:
                entry = decision.entry if isinstance(decision.entry, (int, float)) else None
                if entry and decision.stop_loss <= entry:
                    errors.append("SHORT: stop_loss should be above entry")

        # Position size validation
        if decision.position_size > self.max_position_size:
            errors.append(f"Position size {decision.position_size}% exceeds max {self.max_position_size}%")
            decision.position_size = self.max_position_size

        if decision.position_size < self.MIN_POSITION_SIZE:
            errors.append(f"Position size {decision.position_size}% below minimum {self.MIN_POSITION_SIZE}%")
            decision.position_size = self.MIN_POSITION_SIZE

        # Set validation status
        decision.validation_errors = errors

        # In strict mode, any error invalidates the decision
        if self.strict_mode:
            decision.is_valid = len(errors) == 0
        else:
            # In non-strict mode, only critical errors invalidate
            critical_errors = [e for e in errors if "missing" in e.lower() and "warning" not in e.lower()]
            decision.is_valid = len(critical_errors) == 0

    def validate_decision(self, decision: TradingDecision) -> bool:
        """
        Public method to validate a decision.

        Can be used to re-validate a decision after modifications.

        Args:
            decision: TradingDecision to validate

        Returns:
            True if decision is valid
        """
        self._validate_decision(decision)
        return decision.is_valid


# ============================================================
# CONVENIENCE FUNCTIONS
# ============================================================

def parse_llm_response(response: str) -> TradingDecision:
    """
    Convenience function to parse an LLM response.

    Args:
        response: Raw LLM response text

    Returns:
        Parsed TradingDecision
    """
    parser = DecisionParser()
    return parser.parse_response(response)


def is_valid_trade_decision(decision: TradingDecision) -> Tuple[bool, List[str]]:
    """
    Check if a trade decision is valid for execution.

    Args:
        decision: TradingDecision to check

    Returns:
        Tuple of (is_valid, list of reasons if invalid)
    """
    if not decision.is_valid:
        return False, decision.validation_errors

    if decision.action != "TRADE":
        return False, [f"Decision is {decision.action}, not TRADE"]

    if decision.direction is None:
        return False, ["No direction specified"]

    if decision.stop_loss is None:
        return False, ["No stop loss specified"]

    return True, []


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

__all__ = [
    'TradingDecision',
    'DecisionParser',
    'parse_llm_response',
    'is_valid_trade_decision',
]
