"""
Integration tests for cross-module interactions.
"""
import pytest
import os
import numpy as np
import tempfile
import json

# Try to import the module; skip tests if not found
try:
    from gru_sac_predictor.src import sac_agent
    import tensorflow as tf # Needed for agent init/load
except ImportError:
    sac_agent = None
    tf = None

@pytest.fixture
def sac_agent_for_integration():
    """Provides a basic SAC agent instance."""
    if sac_agent is None or tf is None:
        pytest.skip("SAC Agent module or TF not found.")
    # Use minimal params for saving/loading tests
    agent = sac_agent.SACTradingAgent(
        state_dim=5, action_dim=1,
        buffer_capacity=100, min_buffer_size=10
    )
    # Build models
    try:
        agent.actor(tf.zeros((1, 5)))
        agent.critic1([tf.zeros((1, 5)), tf.zeros((1, 1))])
        agent.critic2([tf.zeros((1, 5)), tf.zeros((1, 1))])
        agent.update_target_networks(tau=1.0)
    except Exception as e:
        pytest.fail(f"Failed to build agent models: {e}")
    return agent

@pytest.mark.skipif(sac_agent is None or tf is None, reason="SAC Agent module or TF not found")
def test_save_load_metadata(sac_agent_for_integration):
    """Test if metadata is saved and loaded correctly."""
    agent = sac_agent_for_integration
    with tempfile.TemporaryDirectory() as tmpdir:
        save_path = os.path.join(tmpdir, "sac_test_save")
        agent.save(save_path)

        # Check if metadata file exists
        meta_path = os.path.join(save_path, 'agent_metadata.json')
        assert os.path.exists(meta_path), "Metadata file was not saved."

        # Create a new agent and load
        new_agent = sac_agent.SACTradingAgent(state_dim=5, action_dim=1)
        loaded_meta = new_agent.load(save_path)

        assert isinstance(loaded_meta, dict), "Load method did not return a dict."
        assert loaded_meta.get('state_dim') == 5, "Loaded state_dim incorrect."
        assert loaded_meta.get('action_dim') == 1, "Loaded action_dim incorrect."
        # Check alpha status (default is auto_tune=True)
        assert loaded_meta.get('log_alpha_saved') == True, "log_alpha status incorrect."

@pytest.mark.skipif(sac_agent is None or tf is None, reason="SAC Agent module or TF not found")
def test_replay_buffer_purge_on_change(sac_agent_for_integration):
    """
    Simulate loading an agent where the edge_threshold has changed
    and verify the buffer is cleared.
    """
    agent_to_save = sac_agent_for_integration
    original_edge_thr = 0.55
    agent_to_save.edge_threshold_config = original_edge_thr # Manually set for saving

    with tempfile.TemporaryDirectory() as tmpdir:
        save_path = os.path.join(tmpdir, "sac_purge_test")

        # 1. Save agent with original threshold in metadata
        agent_to_save.save(save_path)
        meta_path = os.path.join(save_path, 'agent_metadata.json')
        assert os.path.exists(meta_path)
        with open(meta_path, 'r') as f:
             saved_meta = json.load(f)
        assert saved_meta.get('edge_threshold_config') == original_edge_thr

        # 2. Create a new agent instance to load into
        new_agent = sac_agent.SACTradingAgent(
            state_dim=5, action_dim=1,
            buffer_capacity=100, min_buffer_size=10
        )
        # Build models for the new agent
        try:
            new_agent.actor(tf.zeros((1, 5)))
            new_agent.critic1([tf.zeros((1, 5)), tf.zeros((1, 1))])
            new_agent.critic2([tf.zeros((1, 5)), tf.zeros((1, 1))])
            new_agent.update_target_networks(tau=1.0)
        except Exception as e:
            pytest.fail(f"Failed to build new agent models: {e}")

        # Add dummy data to the *new* agent's buffer *before* loading
        for _ in range(20):
            dummy_state = np.random.rand(5).astype(np.float32)
            dummy_action = np.random.rand(1).astype(np.float32)
            new_agent.buffer.add(dummy_state, dummy_action, 0.0, dummy_state, 0.0)
        assert len(new_agent.buffer) == 20, "Buffer should have data before load."

        # 3. Simulate loading with a *different* current edge threshold config
        current_config_edge_thr = 0.60
        assert abs(current_config_edge_thr - original_edge_thr) > 1e-6

        loaded_meta = new_agent.load(save_path)
        saved_edge_thr = loaded_meta.get('edge_threshold_config')

        # 4. Perform the check and clear if needed (simulating pipeline logic)
        if saved_edge_thr is not None and abs(saved_edge_thr - current_config_edge_thr) > 1e-6:
            print(f"\nEdge threshold mismatch detected (Saved={saved_edge_thr}, Current={current_config_edge_thr}). Clearing buffer.")
            new_agent.clear_buffer()
        else:
            print(f"\nEdge threshold match or not saved. Buffer not cleared.")

        # 5. Assert buffer is now empty
        assert len(new_agent.buffer) == 0, "Buffer was not cleared after edge threshold mismatch." 