import pytest
import pandas as pd
import numpy as np
from omegaconf import OmegaConf, DictConfig
from unittest.mock import MagicMock, call # For mocking IOManager and logger
import os
import tempfile
import json

# Adjust the import path based on your project structure
from gru_sac_predictor.src.data_loader import (
    find_missing_bars, 
    _consecutive_gaps, 
    fill_missing_bars,
    report_missing # Import report_missing if we want to test its side effects
)
from gru_sac_predictor.src.io_manager import IOManager # Adjust path if needed

# --- Test Fixtures ---

@pytest.fixture
def gappy_dataframe():
    """Creates a DataFrame with missing timestamps."""
    dates = pd.to_datetime([
        '2023-01-01 00:00:00', '2023-01-01 00:01:00', 
        # Missing 00:02:00, 00:03:00 (2 bars)
        '2023-01-01 00:04:00', 
        # Missing 00:05:00 (1 bar)
        '2023-01-01 00:06:00', '2023-01-01 00:07:00',
        # Missing 00:08:00, 00:09:00, 00:10:00 (3 bars)
        '2023-01-01 00:11:00', '2023-01-01 00:12:00'
    ])
    data = {
        'open': [100, 101, 104, 106, 107, 111, 112],
        'high': [100.5, 101.5, 104.5, 106.5, 107.5, 111.5, 112.5],
        'low': [99.5, 100.5, 103.5, 105.5, 106.5, 110.5, 111.5],
        'close': [100.2, 101.2, 104.2, 106.2, 107.2, 111.2, 112.2],
        'volume': [10, 11, 14, 16, 17, 21, 22]
    }
    df = pd.DataFrame(data, index=dates)
    df.index.name = 'timestamp'
    return df

@pytest.fixture
def base_config():
    """Creates a base OmegaConf config for testing."""
    conf = OmegaConf.create({
        'data': {
            'bar_frequency': '1T', # Using 'T' for minute frequency
            'missing': {
                'strategy': 'neutral', # Default strategy
                'max_gap': 5,
                'interpolate': {
                    'method': 'linear',
                    'limit': 10
                }
            }
        }
        # Add other necessary sections if fill_missing_bars requires them
    })
    return conf

@pytest.fixture
def mock_io_manager():
    """Creates a mock IOManager that saves to a temporary directory."""
    with tempfile.TemporaryDirectory() as tmpdir:
        # Mock the IOManager methods needed by report_missing
        mock_io = MagicMock(spec=IOManager)
        mock_io.results_dir = tmpdir
        
        # Define how save_json should behave
        saved_jsons = {}
        def mock_save_json(data, filename, **kwargs):
            filepath = os.path.join(tmpdir, filename)
            saved_jsons[filename] = data # Store saved data for inspection
            with open(filepath, 'w') as f:
                json.dump(data, f, **kwargs)
            print(f"Mock saving {filename} to {filepath}")
        
        mock_io.save_json.side_effect = mock_save_json
        mock_io.get_artifact_path.side_effect = lambda filename: os.path.join(tmpdir, filename)
        
        # Attach the saved data dictionary for test inspection
        mock_io._saved_jsons = saved_jsons 
        
        yield mock_io # Provide the mock object to the test


@pytest.fixture
def mock_logger():
    """Creates a mock logger."""
    return MagicMock()

# --- Test Functions ---

def test_find_missing_bars(gappy_dataframe):
    expected_missing = pd.to_datetime([
        '2023-01-01 00:02:00', '2023-01-01 00:03:00', 
        '2023-01-01 00:05:00', 
        '2023-01-01 00:08:00', '2023-01-01 00:09:00', '2023-01-01 00:10:00'
    ])
    missing = find_missing_bars(gappy_dataframe, freq='T')
    pd.testing.assert_index_equal(missing, expected_missing, check_names=False)

def test_consecutive_gaps(gappy_dataframe):
    missing = find_missing_bars(gappy_dataframe, freq='T')
    gaps = _consecutive_gaps(missing)
    assert sorted(gaps) == sorted([2, 1, 3]) # Check counts of consecutive missing bars

def test_fill_missing_bars_strategy_drop(gappy_dataframe, base_config, mock_io_manager, mock_logger):
    cfg = base_config.copy()
    cfg.data.missing.strategy = 'drop'
    
    original_df = gappy_dataframe.copy()
    filled_df = fill_missing_bars(original_df, cfg, mock_io_manager, mock_logger)

    # Strategy 'drop' should return the original data, but with imputed flag
    assert 'bar_imputed' in filled_df.columns
    pd.testing.assert_frame_equal(filled_df.drop(columns=['bar_imputed']), original_df)
    
    # Check imputed flags (should be false for existing bars)
    assert not filled_df['bar_imputed'].any()

def test_fill_missing_bars_strategy_neutral(gappy_dataframe, base_config, mock_io_manager, mock_logger):
    cfg = base_config.copy()
    cfg.data.missing.strategy = 'neutral'
    
    filled_df = fill_missing_bars(gappy_dataframe.copy(), cfg, mock_io_manager, mock_logger)
    
    assert 'bar_imputed' in filled_df.columns
    assert len(filled_df) == 13 # Original 7 + 6 missing
    
    # Check imputed timestamps
    missing_ts = find_missing_bars(gappy_dataframe, freq='T')
    pd.testing.assert_index_equal(filled_df[filled_df['bar_imputed']].index, missing_ts)
    assert not filled_df[~filled_df['bar_imputed']].index.isin(missing_ts).any()

    # Check neutral fill logic for a specific imputed bar (00:02:00)
    imputed_bar = filled_df.loc['2023-01-01 00:02:00']
    last_close = gappy_dataframe.loc['2023-01-01 00:01:00', 'close']
    assert imputed_bar['open'] == last_close
    assert imputed_bar['high'] == last_close
    assert imputed_bar['low'] == last_close
    assert imputed_bar['close'] == last_close
    assert imputed_bar['volume'] == 0
    assert imputed_bar['bar_imputed'] == True

    # Check a non-imputed bar remains unchanged
    original_bar = gappy_dataframe.loc['2023-01-01 00:04:00']
    filled_bar = filled_df.loc['2023-01-01 00:04:00']
    assert filled_bar['close'] == original_bar['close']
    assert filled_bar['volume'] > 0 # Should not be 0
    assert filled_bar['bar_imputed'] == False
    
    # Check log message (optional, depends on mocking detail)
    mock_logger.warning.assert_called_once()
    assert 'Detected 6 missing bars' in mock_logger.warning.call_args[0][0]
    # Check saved report (optional)
    assert 'missing_bars_summary.json' in mock_io_manager._saved_jsons
    report_data = mock_io_manager._saved_jsons['missing_bars_summary.json']
    assert report_data['total_missing_bars'] == 6
    assert report_data['longest_consecutive_gap'] == 3
    assert report_data['applied_strategy'] == 'neutral'


def test_fill_missing_bars_strategy_ffill(gappy_dataframe, base_config, mock_io_manager, mock_logger):
    cfg = base_config.copy()
    cfg.data.missing.strategy = 'ffill'
    
    filled_df = fill_missing_bars(gappy_dataframe.copy(), cfg, mock_io_manager, mock_logger)
    
    assert 'bar_imputed' in filled_df.columns
    assert len(filled_df) == 13
    
    # Check imputed timestamps
    missing_ts = find_missing_bars(gappy_dataframe, freq='T')
    pd.testing.assert_index_equal(filled_df[filled_df['bar_imputed']].index, missing_ts)

    # Check ffill logic for a specific imputed bar (00:02:00)
    imputed_bar = filled_df.loc['2023-01-01 00:02:00']
    prev_bar = gappy_dataframe.loc['2023-01-01 00:01:00']
    assert imputed_bar['open'] == prev_bar['open']
    assert imputed_bar['high'] == prev_bar['high']
    assert imputed_bar['low'] == prev_bar['low']
    assert imputed_bar['close'] == prev_bar['close']
    assert imputed_bar['volume'] == prev_bar['volume']
    assert imputed_bar['bar_imputed'] == True

    # Check another imputed bar (00:05:00)
    imputed_bar_2 = filled_df.loc['2023-01-01 00:05:00']
    prev_bar_2 = gappy_dataframe.loc['2023-01-01 00:04:00']
    assert imputed_bar_2['close'] == prev_bar_2['close']
    assert imputed_bar_2['bar_imputed'] == True


def test_fill_missing_bars_strategy_interpolate(gappy_dataframe, base_config, mock_io_manager, mock_logger):
    cfg = base_config.copy()
    cfg.data.missing.strategy = 'interpolate'
    cfg.data.missing.interpolate.method = 'linear' 
    
    filled_df = fill_missing_bars(gappy_dataframe.copy(), cfg, mock_io_manager, mock_logger)
    
    assert 'bar_imputed' in filled_df.columns
    assert len(filled_df) == 13
    
    # Check imputed timestamps
    missing_ts = find_missing_bars(gappy_dataframe, freq='T')
    pd.testing.assert_index_equal(filled_df[filled_df['bar_imputed']].index, missing_ts)

    # Check interpolation logic for close price at 00:02:00 and 00:03:00
    # These are between 00:01:00 (101.2) and 00:04:00 (104.2) - 3 steps total
    close_01 = gappy_dataframe.loc['2023-01-01 00:01:00', 'close']
    close_04 = gappy_dataframe.loc['2023-01-01 00:04:00', 'close']
    expected_close_02 = close_01 + (close_04 - close_01) / 3.0 * 1
    expected_close_03 = close_01 + (close_04 - close_01) / 3.0 * 2
    
    assert np.isclose(filled_df.loc['2023-01-01 00:02:00', 'close'], expected_close_02)
    assert np.isclose(filled_df.loc['2023-01-01 00:03:00', 'close'], expected_close_03)
    assert filled_df.loc['2023-01-01 00:02:00', 'bar_imputed'] == True
    assert filled_df.loc['2023-01-01 00:03:00', 'bar_imputed'] == True

    # Check interpolation logic for volume at 00:05:00
    # This is between 00:04:00 (14) and 00:06:00 (16) - 2 steps total
    vol_04 = gappy_dataframe.loc['2023-01-01 00:04:00', 'volume']
    vol_06 = gappy_dataframe.loc['2023-01-01 00:06:00', 'volume']
    expected_vol_05 = vol_04 + (vol_06 - vol_04) / 2.0 * 1
    assert np.isclose(filled_df.loc['2023-01-01 00:05:00', 'volume'], expected_vol_05)
    assert filled_df.loc['2023-01-01 00:05:00', 'bar_imputed'] == True

def test_fill_missing_bars_max_gap_exceeded(gappy_dataframe, base_config, mock_io_manager, mock_logger):
    cfg = base_config.copy()
    cfg.data.missing.max_gap = 2 # Set max gap lower than the actual max gap (3)
    
    with pytest.raises(ValueError) as excinfo:
        fill_missing_bars(gappy_dataframe.copy(), cfg, mock_io_manager, mock_logger)
    
    assert "Longest consecutive gap (3) exceeds maximum allowed (2)" in str(excinfo.value)
    # Verify report was still generated before the error
    mock_logger.warning.assert_called_once() 
    assert 'missing_bars_summary.json' in mock_io_manager._saved_jsons

def test_fill_missing_bars_no_missing_data(gappy_dataframe, base_config, mock_io_manager, mock_logger):
    # Create a dataframe with no gaps
    no_gap_df = pd.DataFrame({
        'open': [1, 2, 3], 'high': [1, 2, 3], 'low': [1, 2, 3], 'close': [1, 2, 3], 'volume': [10, 20, 30]
    }, index=pd.date_range('2023-01-01', periods=3, freq='T'))
    
    filled_df = fill_missing_bars(no_gap_df.copy(), base_config, mock_io_manager, mock_logger)
    
    assert 'bar_imputed' in filled_df.columns
    assert not filled_df['bar_imputed'].any() # No bars should be marked as imputed
    pd.testing.assert_frame_equal(filled_df.drop(columns=['bar_imputed']), no_gap_df)
    mock_logger.info.assert_any_call("No missing bars detected.") # Check info log 