"""Tests for DCA execution engine."""

import numpy as np
import pytest
from src.composed.dca_executor import (
    DCAParams, DCAOrder, DCAResult,
    create_dca_order, compute_chunk_size, execute_chunk, simulate_dca,
)
from src.composed.myco_surface import MycoSystem, MycoSystemConfig
from src.primitives.twap_oracle import (
    TWAPOracleParams, create_oracle, record_observation,
)


def _bootstrap_system(n_assets=3, amount=10000.0):
    """Helper: create and bootstrap a MycoSystem."""
    config = MycoSystemConfig(n_reserve_assets=n_assets)
    system = MycoSystem(config)
    bootstrap = np.full(n_assets, amount / n_assets)
    system.deposit(bootstrap, 0.0)
    return system


class TestFixedDCA:
    def test_equal_chunks(self):
        """Fixed strategy splits into equal chunks."""
        params = DCAParams(total_amount=1000.0, n_chunks=10, interval=1.0)
        order = create_dca_order(params, start_time=1.0)
        oracle = create_oracle()
        chunk = compute_chunk_size(order, oracle)
        assert chunk == pytest.approx(100.0, abs=0.01)

    def test_budget_exhaustion(self):
        """All chunks together should spend the full budget."""
        system = _bootstrap_system()
        params = DCAParams(total_amount=1000.0, n_chunks=5, interval=1.0)
        order = create_dca_order(params, start_time=1.0)
        oracle = create_oracle()

        for i in range(5):
            t = 1.0 + i * 1.0
            system, order, oracle, tokens = execute_chunk(system, order, oracle, t)

        assert order.is_complete
        assert order.total_spent == pytest.approx(1000.0, abs=0.01)
        assert order.chunks_executed == 5

    def test_token_tracking(self):
        """Total tokens received should be positive and match history."""
        system = _bootstrap_system()
        params = DCAParams(total_amount=500.0, n_chunks=5, interval=1.0)
        order = create_dca_order(params, start_time=1.0)
        oracle = create_oracle()

        total_from_chunks = 0.0
        for i in range(5):
            t = 1.0 + i * 1.0
            system, order, oracle, tokens = execute_chunk(system, order, oracle, t)
            total_from_chunks += tokens

        assert order.total_tokens_received > 0
        assert order.total_tokens_received == pytest.approx(total_from_chunks, rel=1e-10)
        assert len(order.history) == 5

    def test_avg_price_computed(self):
        """Average price should be total_spent / total_tokens."""
        system = _bootstrap_system()
        params = DCAParams(total_amount=1000.0, n_chunks=4, interval=1.0)
        order = create_dca_order(params, start_time=1.0)
        oracle = create_oracle()

        for i in range(4):
            system, order, oracle, _ = execute_chunk(
                system, order, oracle, 1.0 + i,
            )

        expected_avg = order.total_spent / order.total_tokens_received
        assert order.avg_price == pytest.approx(expected_avg, rel=1e-10)

    def test_no_execution_after_complete(self):
        """Executing on a complete order returns 0 tokens."""
        system = _bootstrap_system()
        params = DCAParams(total_amount=100.0, n_chunks=1, interval=1.0)
        order = create_dca_order(params, start_time=1.0)
        oracle = create_oracle()

        system, order, oracle, t1 = execute_chunk(system, order, oracle, 1.0)
        assert order.is_complete

        system, order, oracle, t2 = execute_chunk(system, order, oracle, 2.0)
        assert t2 == 0.0


class TestTWAPAwareDCA:
    def _make_oracle_with_history(self, base_price=1.0, n=10):
        """Create oracle with price history at base_price."""
        oracle = create_oracle(TWAPOracleParams(default_window=100.0))
        for t in range(n):
            oracle = record_observation(oracle, base_price, float(t))
        return oracle

    def test_buys_more_when_cheap(self):
        """When spot < TWAP, chunk should be larger than base."""
        # Oracle has TWAP ~ 10.0
        oracle = self._make_oracle_with_history(base_price=10.0)
        # Spot drops to 8.0
        oracle = record_observation(oracle, 8.0, 11.0)

        params = DCAParams(
            total_amount=1000.0, n_chunks=10, interval=1.0,
            strategy="twap_aware", max_deviation=0.3,
        )
        order = create_dca_order(params, start_time=12.0)
        chunk = compute_chunk_size(order, oracle)
        base = 1000.0 / 10
        assert chunk > base

    def test_buys_less_when_expensive(self):
        """When spot > TWAP, chunk should be smaller than base."""
        oracle = self._make_oracle_with_history(base_price=10.0)
        # Spot rises to 12.0
        oracle = record_observation(oracle, 12.0, 11.0)

        params = DCAParams(
            total_amount=1000.0, n_chunks=10, interval=1.0,
            strategy="twap_aware", max_deviation=0.3,
        )
        order = create_dca_order(params, start_time=12.0)
        chunk = compute_chunk_size(order, oracle)
        base = 1000.0 / 10
        assert chunk < base

    def test_equal_when_spot_equals_twap(self):
        """When spot == TWAP, chunk == base."""
        oracle = self._make_oracle_with_history(base_price=10.0)

        params = DCAParams(
            total_amount=1000.0, n_chunks=10, interval=1.0,
            strategy="twap_aware", max_deviation=0.3,
        )
        order = create_dca_order(params, start_time=11.0)
        chunk = compute_chunk_size(order, oracle)
        base = 1000.0 / 10
        assert chunk == pytest.approx(base, rel=0.05)

    def test_chunk_clamped_to_remaining(self):
        """Chunk size never exceeds remaining budget."""
        oracle = self._make_oracle_with_history(base_price=10.0)
        oracle = record_observation(oracle, 5.0, 11.0)  # Very cheap

        params = DCAParams(
            total_amount=100.0, n_chunks=10, interval=1.0,
            strategy="twap_aware", max_deviation=0.3,
        )
        order = create_dca_order(params, start_time=12.0)
        order.total_spent = 95.0  # Only 5 remaining
        chunk = compute_chunk_size(order, oracle)
        assert chunk <= 5.0


class TestDCAvsLumpSum:
    def test_simulation_runs(self):
        """Basic DCA simulation completes without error."""
        config = MycoSystemConfig(n_reserve_assets=3)
        params = DCAParams(
            total_amount=5000.0, n_chunks=10, interval=1.0,
        )
        result = simulate_dca(config, params)

        assert result.order.is_complete
        assert result.order.total_tokens_received > 0
        assert result.lump_sum_tokens > 0

    def test_dca_advantage_computed(self):
        """DCA advantage is a finite number."""
        config = MycoSystemConfig(n_reserve_assets=3)
        params = DCAParams(
            total_amount=5000.0, n_chunks=10, interval=1.0,
        )
        result = simulate_dca(config, params)
        assert np.isfinite(result.dca_advantage)

    def test_twap_aware_simulation(self):
        """TWAP-aware strategy also completes."""
        config = MycoSystemConfig(n_reserve_assets=3)
        params = DCAParams(
            total_amount=5000.0, n_chunks=10, interval=1.0,
            strategy="twap_aware", twap_window=50.0,
        )
        result = simulate_dca(config, params)
        assert result.order.is_complete
        assert result.order.total_tokens_received > 0


class TestEdgeCases:
    def test_single_chunk_equals_lump_sum(self):
        """With n_chunks=1, DCA is equivalent to lump sum."""
        config = MycoSystemConfig(n_reserve_assets=3)
        params_dca = DCAParams(
            total_amount=5000.0, n_chunks=1, interval=1.0,
        )
        result = simulate_dca(config, params_dca)

        # With one chunk, DCA and lump sum should get similar tokens
        # (not exact due to different bootstrap timing)
        ratio = result.order.total_tokens_received / result.lump_sum_tokens
        assert 0.95 < ratio < 1.05

    def test_zero_bootstrap(self):
        """DCA on a fresh system (bootstrap via first chunk)."""
        config = MycoSystemConfig(n_reserve_assets=3)
        params = DCAParams(
            total_amount=1000.0, n_chunks=5, interval=1.0,
        )
        # Small bootstrap so system exists
        result = simulate_dca(config, params, bootstrap_amount=100.0)
        assert result.order.is_complete
        assert result.order.total_tokens_received > 0

    def test_large_dca_relative_to_system(self):
        """DCA for an amount much larger than bootstrap."""
        config = MycoSystemConfig(n_reserve_assets=3)
        params = DCAParams(
            total_amount=50000.0, n_chunks=20, interval=1.0,
        )
        result = simulate_dca(config, params, bootstrap_amount=1000.0)
        assert result.order.is_complete
        assert result.order.total_spent == pytest.approx(50000.0, abs=1.0)