"""
Tests for the strategy generation pipeline.
Tests the flow: AI description → Blockly XML → JSON → Python
"""
import json
import pytest
import subprocess
import xml.etree.ElementTree as ET
import sys
import os

# Add src to path for imports
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))

from PythonGenerator import PythonGenerator


class TestStrategyBuilder:
    """Tests for the CmdForge strategy-builder tool."""

    @staticmethod
    def _get_blocks(root):
        """Get all block elements, handling XML namespaces."""
        # Blockly uses namespace https://developers.google.com/blockly/xml
        # Elements may be prefixed with {namespace}block or just block
        blocks = []
        for elem in root.iter():
            # Get local name without namespace
            tag = elem.tag.split('}')[-1] if '}' in elem.tag else elem.tag
            if tag == 'block':
                blocks.append(elem)
        return blocks

    @pytest.mark.integration
    def test_simple_rsi_strategy(self):
        """Test generating a simple RSI-based strategy."""
        input_data = {
            "description": "Buy when RSI is below 30 and sell when RSI is above 70",
            "indicators": [{"name": "RSI", "outputs": ["RSI"]}],
            "signals": [],
            "default_source": {"exchange": "binance", "market": "BTC/USDT", "timeframe": "5m"}
        }

        result = subprocess.run(
            ['strategy-builder'],
            input=json.dumps(input_data),
            capture_output=True,
            text=True,
            timeout=120
        )

        assert result.returncode == 0, f"Strategy builder failed: {result.stderr}"

        # Validate XML structure
        xml_output = result.stdout.strip()
        root = ET.fromstring(xml_output)

        # Check it has the required structure
        root_tag = root.tag.split('}')[-1] if '}' in root.tag else root.tag
        assert root_tag == 'xml', f"Root should be 'xml', got {root_tag}"

        blocks = self._get_blocks(root)
        assert len(blocks) >= 2, f"Should have at least 2 blocks (buy and sell), got {len(blocks)}"

        # Check for execute_if blocks
        execute_if_blocks = [b for b in blocks if b.get('type') == 'execute_if']
        assert len(execute_if_blocks) >= 2, "Should have at least 2 execute_if blocks"

        # Check for trade_action blocks
        trade_action_blocks = [b for b in blocks if b.get('type') == 'trade_action']
        assert len(trade_action_blocks) >= 2, "Should have buy and sell trade actions"

        # Check for indicator blocks
        indicator_blocks = [b for b in blocks if 'indicator_RSI' in (b.get('type') or '')]
        assert len(indicator_blocks) >= 2, "Should use RSI indicator"

    @pytest.mark.integration
    def test_ema_crossover_strategy(self):
        """Test generating an EMA crossover strategy."""
        input_data = {
            "description": "Buy when EMA 20 crosses above EMA 50, sell when EMA 20 crosses below EMA 50",
            "indicators": [
                {"name": "EMA_20", "outputs": ["ema"]},
                {"name": "EMA_50", "outputs": ["ema"]}
            ],
            "signals": [],
            "default_source": {"exchange": "binance", "market": "ETH/USDT", "timeframe": "1h"}
        }

        result = subprocess.run(
            ['strategy-builder'],
            input=json.dumps(input_data),
            capture_output=True,
            text=True,
            timeout=120
        )

        assert result.returncode == 0, f"Strategy builder failed: {result.stderr}"

        xml_output = result.stdout.strip()
        root = ET.fromstring(xml_output)

        # Check for EMA indicator blocks
        blocks = self._get_blocks(root)
        ema_20_blocks = [b for b in blocks if 'indicator_EMA_20' in (b.get('type') or '')]
        ema_50_blocks = [b for b in blocks if 'indicator_EMA_50' in (b.get('type') or '')]

        assert len(ema_20_blocks) >= 1, "Should use EMA_20 indicator"
        assert len(ema_50_blocks) >= 1, "Should use EMA_50 indicator"

    @pytest.mark.integration
    def test_no_indicators_price_only(self):
        """Test generating a price-based strategy without indicators."""
        input_data = {
            "description": "Buy when price drops 5% from previous candle, sell when price rises 3%",
            "indicators": [],
            "signals": [],
            "default_source": {"exchange": "binance", "market": "BTC/USDT", "timeframe": "5m"}
        }

        result = subprocess.run(
            ['strategy-builder'],
            input=json.dumps(input_data),
            capture_output=True,
            text=True,
            timeout=120
        )

        assert result.returncode == 0, f"Strategy builder failed: {result.stderr}"

        xml_output = result.stdout.strip()
        root = ET.fromstring(xml_output)

        # Should be valid XML
        assert root is not None

    @pytest.mark.integration
    def test_missing_indicator_error(self):
        """Test that strategy mentioning indicators without config fails."""
        input_data = {
            "description": "Buy when RSI is below 30",
            "indicators": [],  # No indicators configured
            "signals": [],
            "default_source": {"exchange": "binance", "market": "BTC/USDT", "timeframe": "5m"}
        }

        result = subprocess.run(
            ['strategy-builder'],
            input=json.dumps(input_data),
            capture_output=True,
            text=True,
            timeout=120
        )

        assert result.returncode != 0, "Should fail when indicators mentioned but not configured"
        assert "indicator" in result.stderr.lower(), "Error should mention indicators"


class TestPythonGenerator:
    """Tests for the PythonGenerator class."""

    def test_simple_execute_if(self):
        """Test generating Python from a simple execute_if block."""
        strategy_json = {
            "type": "strategy",
            "statements": [
                {
                    "type": "execute_if",
                    "inputs": {
                        "CONDITION": {
                            "type": "comparison",
                            "operator": ">",
                            "inputs": {
                                "LEFT": {"type": "current_price"},
                                "RIGHT": {"type": "dynamic_value", "values": [50000]}
                            }
                        }
                    },
                    "statements": {
                        "DO": [
                            {
                                "type": "trade_action",
                                "trade_type": "buy",
                                "inputs": {"size": 0.01}
                            }
                        ]
                    }
                }
            ]
        }

        generator = PythonGenerator(
            default_source={"exchange": "binance", "market": "BTC/USDT", "timeframe": "5m"},
            strategy_id="test_strategy"
        )

        result = generator.generate(strategy_json)

        assert "generated_code" in result
        code = result["generated_code"]

        # Check for expected code elements
        assert "def next():" in code
        assert "if " in code
        assert "get_current_price" in code
        assert "trade_order" in code

    def test_indicator_condition(self):
        """Test generating Python with indicator conditions."""
        strategy_json = {
            "type": "strategy",
            "statements": [
                {
                    "type": "execute_if",
                    "inputs": {
                        "CONDITION": {
                            "type": "comparison",
                            "operator": "<",
                            "inputs": {
                                "LEFT": {
                                    "type": "indicator_RSI",
                                    "fields": {"OUTPUT": "RSI"}
                                },
                                "RIGHT": {"type": "dynamic_value", "values": [30]}
                            }
                        }
                    },
                    "statements": {
                        "DO": [
                            {
                                "type": "trade_action",
                                "trade_type": "buy",
                                "inputs": {"size": 0.1}
                            }
                        ]
                    }
                }
            ]
        }

        generator = PythonGenerator(
            default_source={"exchange": "binance", "market": "BTC/USDT", "timeframe": "5m"},
            strategy_id="test_indicator"
        )

        result = generator.generate(strategy_json)
        code = result["generated_code"]

        # Check for indicator processing
        assert "process_indicator" in code
        assert "RSI" in code

        # Check indicators are tracked
        assert len(result["indicators"]) > 0

    def test_logical_and_condition(self):
        """Test generating Python with logical AND conditions."""
        strategy_json = {
            "type": "strategy",
            "statements": [
                {
                    "type": "execute_if",
                    "inputs": {
                        "CONDITION": {
                            "type": "logical_and",
                            "inputs": {
                                "left": {
                                    "type": "comparison",
                                    "operator": "<",
                                    "inputs": {
                                        "LEFT": {"type": "indicator_RSI", "fields": {"OUTPUT": "RSI"}},
                                        "RIGHT": {"type": "dynamic_value", "values": [30]}
                                    }
                                },
                                "right": {
                                    "type": "flag_is_set",
                                    "flag_name": "bought",
                                    "flag_value": False
                                }
                            }
                        }
                    },
                    "statements": {
                        "DO": [
                            {"type": "trade_action", "trade_type": "buy", "inputs": {"size": 0.01}}
                        ]
                    }
                }
            ]
        }

        generator = PythonGenerator(
            default_source={"exchange": "binance", "market": "BTC/USDT", "timeframe": "5m"},
            strategy_id="test_and"
        )

        result = generator.generate(strategy_json)
        code = result["generated_code"]

        # Check for logical AND
        assert " and " in code
        assert "flags.get" in code

    def test_set_flag(self):
        """Test generating Python for flag setting."""
        strategy_json = {
            "type": "strategy",
            "statements": [
                {
                    "type": "set_flag",
                    "flag_name": "bought",
                    "flag_value": "True"
                }
            ]
        }

        generator = PythonGenerator(
            default_source={"exchange": "binance", "market": "BTC/USDT", "timeframe": "5m"},
            strategy_id="test_flag"
        )

        result = generator.generate(strategy_json)
        code = result["generated_code"]

        # Check for flag setting
        assert "flags['bought']" in code
        assert "True" in code

        # Check flag is tracked
        assert "bought" in result["flags_used"]

    def test_trade_action_with_options(self):
        """Test generating Python for trade with stop loss and take profit."""
        strategy_json = {
            "type": "strategy",
            "statements": [
                {
                    "type": "trade_action",
                    "trade_type": "buy",
                    "inputs": {"size": 0.1},
                    "trade_options": [
                        {"type": "stop_loss", "inputs": {"stop_loss": 45000}},
                        {"type": "take_profit", "inputs": {"take_profit": 55000}}
                    ]
                }
            ]
        }

        generator = PythonGenerator(
            default_source={"exchange": "binance", "market": "BTC/USDT", "timeframe": "5m"},
            strategy_id="test_trade_options"
        )

        result = generator.generate(strategy_json)
        code = result["generated_code"]

        # Check for trade order call
        assert "trade_order" in code
        assert "buy" in code

    def test_math_operation(self):
        """Test generating Python for math operations."""
        strategy_json = {
            "type": "strategy",
            "statements": [
                {
                    "type": "execute_if",
                    "inputs": {
                        "CONDITION": {
                            "type": "comparison",
                            "operator": ">",
                            "inputs": {
                                "LEFT": {"type": "current_price"},
                                "RIGHT": {
                                    "type": "math_operation",
                                    "inputs": {
                                        "operator": "MULTIPLY",
                                        "left_operand": {"type": "indicator_SMA", "fields": {"OUTPUT": "sma"}},
                                        "right_operand": {"type": "dynamic_value", "values": [1.02]}
                                    }
                                }
                            }
                        }
                    },
                    "statements": {
                        "DO": [
                            {"type": "trade_action", "trade_type": "sell", "inputs": {"size": 0.01}}
                        ]
                    }
                }
            ]
        }

        generator = PythonGenerator(
            default_source={"exchange": "binance", "market": "BTC/USDT", "timeframe": "5m"},
            strategy_id="test_math"
        )

        result = generator.generate(strategy_json)
        code = result["generated_code"]

        # Check for math operation
        assert "*" in code  # Multiply operator