from __future__ import annotations

import json

import numpy as np
import pandas as pd
import pytest

from multi_swarm.backtest.orders import Side
from multi_swarm.protocol.compiler import compile_strategy
from multi_swarm.protocol.parser import parse_strategy


@pytest.fixture
def ohlcv() -> pd.DataFrame:
    idx = pd.date_range("2024-01-01", periods=200, freq="1h", tz="UTC")
    close = np.linspace(100, 120, 200)
    return pd.DataFrame(
        {
            "open": close,
            "high": close + 0.5,
            "low": close - 0.5,
            "close": close,
            "volume": 1.0,
        },
        index=idx,
    )


def test_compile_simple_long(ohlcv: pd.DataFrame) -> None:
    src = json.dumps(
        {
            "rules": [
                {
                    "condition": {
                        "op": "lt",
                        "args": [
                            {"kind": "indicator", "name": "rsi", "params": [14]},
                            {"kind": "literal", "value": 100.0},
                        ],
                    },
                    "action": "entry-long",
                }
            ]
        }
    )
    ast = parse_strategy(src)
    fn = compile_strategy(ast)
    signals = fn(ohlcv)
    assert isinstance(signals, pd.Series)
    assert (signals == Side.LONG).all() or (signals.dropna() == Side.LONG).all()


def test_compile_no_match_is_flat(ohlcv: pd.DataFrame) -> None:
    src = json.dumps(
        {
            "rules": [
                {
                    "condition": {
                        "op": "gt",
                        "args": [
                            {"kind": "indicator", "name": "rsi", "params": [14]},
                            {"kind": "literal", "value": 1000.0},
                        ],
                    },
                    "action": "entry-long",
                }
            ]
        }
    )
    ast = parse_strategy(src)
    fn = compile_strategy(ast)
    signals = fn(ohlcv)
    assert (signals == Side.FLAT).any()


def test_compile_two_rules_priority(ohlcv: pd.DataFrame) -> None:
    src = json.dumps(
        {
            "rules": [
                {
                    "condition": {
                        "op": "gt",
                        "args": [
                            {"kind": "feature", "name": "close"},
                            {"kind": "literal", "value": 110.0},
                        ],
                    },
                    "action": "entry-long",
                },
                {
                    "condition": {
                        "op": "lt",
                        "args": [
                            {"kind": "feature", "name": "close"},
                            {"kind": "literal", "value": 105.0},
                        ],
                    },
                    "action": "entry-short",
                },
            ]
        }
    )
    ast = parse_strategy(src)
    fn = compile_strategy(ast)
    signals = fn(ohlcv)
    last = signals.iloc[-1]
    assert last == Side.LONG  # close finale e' 120, regola 1 matcha


def test_compile_hour_feature_returns_index_hour(ohlcv: pd.DataFrame) -> None:
    src = json.dumps(
        {
            "rules": [
                {
                    "condition": {
                        "op": "gt",
                        "args": [
                            {"kind": "feature", "name": "hour"},
                            {"kind": "literal", "value": -1.0},
                        ],
                    },
                    "action": "entry-long",
                }
            ]
        }
    )
    ast = parse_strategy(src)
    fn = compile_strategy(ast)
    signal = fn(ohlcv)
    # All rows have hour >= 0 > -1, so all entry-long.
    assert (signal == Side.LONG).all()


def test_compile_dow_feature_monday_is_zero(ohlcv: pd.DataFrame) -> None:
    # 2024-01-01 is Monday -> dow=0; eq(dow, 0) gates LONG on Monday rows only.
    src = json.dumps(
        {
            "rules": [
                {
                    "condition": {
                        "op": "eq",
                        "args": [
                            {"kind": "feature", "name": "dow"},
                            {"kind": "literal", "value": 0.0},
                        ],
                    },
                    "action": "entry-long",
                }
            ]
        }
    )
    ast = parse_strategy(src)
    fn = compile_strategy(ast)
    signal = fn(ohlcv)
    monday_rows = signal[signal.index.dayofweek == 0]
    other_rows = signal[signal.index.dayofweek != 0]
    assert (monday_rows == Side.LONG).all()
    assert (other_rows == Side.FLAT).all()


def test_compile_is_weekend_returns_zero_one(ohlcv: pd.DataFrame) -> None:
    src = json.dumps(
        {
            "rules": [
                {
                    "condition": {
                        "op": "eq",
                        "args": [
                            {"kind": "feature", "name": "is_weekend"},
                            {"kind": "literal", "value": 1.0},
                        ],
                    },
                    "action": "entry-long",
                }
            ]
        }
    )
    ast = parse_strategy(src)
    fn = compile_strategy(ast)
    signal = fn(ohlcv)
    weekend = signal[signal.index.dayofweek >= 5]
    weekdays = signal[signal.index.dayofweek < 5]
    assert (weekend == Side.LONG).all()
    assert (weekdays == Side.FLAT).all()