"""
Tests for CLI functionality.
"""

import pytest
import tempfile
import shutil
import yaml
from pathlib import Path
from unittest.mock import patch, MagicMock
import pandas as pd
import numpy as np

from medtda.cli import (
    create_parser,
    load_config,
    merge_config_with_args,
    validate_args,
    detect_mode,
    save_config,
    main
)


@pytest.fixture
def temp_dir():
    """Create a temporary directory for test outputs."""
    tmp = tempfile.mkdtemp()
    yield Path(tmp)
    shutil.rmtree(tmp)


@pytest.fixture
def sample_config_dict():
    """Sample configuration dictionary."""
    return {
        'preprocessing': {
            'normalize': True,
            'normalize_method': 'zscore',
            'spacing': [1.0, 1.0, 2.0],
            'window': [50, 100],
            'crop_roi': True,
            'roi_padding': 2,
            'background_value': 0,
            'label': 2,
        },
        'persistent_homology': {
            'filtration': 'superlevel',
            'construction': 'V',
            'max_dimension': 1,
        },
        'vectorization': {
            'methods': ['persistence_stats', 'betti_curve'],
            'save_barcodes': True,
        },
        'parallel': {
            'workers': 4,
        },
        'output': {
            'output_dir': './test_output',
            'verbose': True,
            'quiet': False,
        }
    }


@pytest.fixture
def sample_config_file(temp_dir, sample_config_dict):
    """Create a sample config YAML file."""
    config_path = temp_dir / "config.yaml"
    with open(config_path, 'w') as f:
        yaml.dump(sample_config_dict, f)
    return config_path


@pytest.fixture
def sample_csv_file(temp_dir):
    """Create a sample CSV file for batch processing."""
    csv_path = temp_dir / "cases.csv"
    df = pd.DataFrame({
        'id': ['case001', 'case002', 'case003'],
        'image_path': ['img1.nii.gz', 'img2.nii.gz', 'img3.nii.gz'],
        'mask_path': ['mask1.nii.gz', '', 'mask3.nii.gz']
    })
    df.to_csv(csv_path, index=False)
    return csv_path


@pytest.fixture
def sample_image_file(temp_dir):
    """Create a dummy image file."""
    img_path = temp_dir / "test_image.nii.gz"
    img_path.touch()
    return img_path


class TestArgumentParser:
    """Test argument parser creation and functionality."""
    
    def test_parser_creation(self):
        """Test that parser is created successfully."""
        parser = create_parser()
        assert parser is not None
        assert parser.prog == 'medtda'
    
    def test_parser_required_args(self, sample_image_file):
        """Test required arguments."""
        parser = create_parser()
        args = parser.parse_args([
            str(sample_image_file),
            '--output-dir', './output'
        ])
        assert args.input == str(sample_image_file)
        assert args.output_dir == './output'
    
    def test_parser_optional_args(self, sample_image_file):
        """Test optional arguments."""
        parser = create_parser()
        args = parser.parse_args([
            str(sample_image_file),
            '--output-dir', './output',
            '--mask', 'mask.nii.gz',
            '--normalize',
            '--normalize-method', 'zscore',
            '--spacing', '1.0', '1.0', '2.0',
            '--window', '50', '100',
            '--filtration', 'superlevel',
            '--construction', 'V',
            '--methods', 'persistence_stats', 'betti_curve',
            '--workers', '4',
            '--save-barcodes',
            '--verbose'
        ])
        
        assert args.mask == 'mask.nii.gz'
        assert args.normalize is True
        assert args.normalize_method == 'zscore'
        assert args.spacing == [1.0, 1.0, 2.0]
        assert args.window == [50.0, 100.0]
        assert args.filtration == 'superlevel'
        assert args.construction == 'V'
        assert args.methods == ['persistence_stats', 'betti_curve']
        assert args.workers == 4
        assert args.save_barcodes is True
        assert args.verbose is True


class TestConfigLoading:
    """Test configuration loading and merging."""
    
    def test_load_valid_config(self, sample_config_file):
        """Test loading a valid config file."""
        config = load_config(str(sample_config_file))
        assert config is not None
        assert 'preprocessing' in config
        assert 'persistent_homology' in config
        assert config['preprocessing']['normalize'] is True
    
    def test_load_nonexistent_config(self):
        """Test loading a non-existent config file."""
        with pytest.raises(ValueError, match="does not exist"):
            load_config('nonexistent_config.yaml')
    
    def test_load_invalid_yaml(self, temp_dir):
        """Test loading an invalid YAML file."""
        invalid_yaml = temp_dir / "invalid.yaml"
        with open(invalid_yaml, 'w') as f:
            f.write("invalid: yaml: content: [")
        
        with pytest.raises(ValueError, match="Invalid YAML"):
            load_config(str(invalid_yaml))
    
    def test_merge_config_with_args(self, sample_config_dict):
        """Test merging config with CLI arguments."""
        parser = create_parser()
        args = parser.parse_args(['image.nii.gz', '--output-dir', './output'])
        
        merged = merge_config_with_args(sample_config_dict, args)
        
        # Config values should be applied
        assert merged.normalize is True
        assert merged.normalize_method == 'zscore'
        assert merged.spacing == [1.0, 1.0, 2.0]
        assert merged.workers == 4
    
    def test_cli_overrides_config(self, sample_config_dict):
        """Test that CLI args override config values."""
        parser = create_parser()
        args = parser.parse_args([
            'image.nii.gz',
            '--output-dir', './cli_output',
            '--workers', '8',
            '--normalize-method', 'minmax'
        ])
        
        merged = merge_config_with_args(sample_config_dict, args)
        
        # CLI values should override config
        assert merged.workers == 8
        assert merged.output_dir == './cli_output'


class TestValidation:
    """Test argument validation."""
    
    def test_validate_missing_output_dir(self, sample_image_file):
        """Test validation fails when output_dir is missing."""
        parser = create_parser()
        args = parser.parse_args([str(sample_image_file)])
        
        with pytest.raises(ValueError, match="--output-dir is required"):
            validate_args(args)
    
    def test_validate_nonexistent_input(self, temp_dir):
        """Test validation fails for non-existent input file."""
        parser = create_parser()
        args = parser.parse_args([
            'nonexistent.nii.gz',
            '--output-dir', str(temp_dir)
        ])
        
        with pytest.raises(ValueError, match="does not exist"):
            validate_args(args)
    
    def test_validate_nonexistent_mask(self, sample_image_file, temp_dir):
        """Test validation fails for non-existent mask file."""
        parser = create_parser()
        args = parser.parse_args([
            str(sample_image_file),
            '--mask', 'nonexistent_mask.nii.gz',
            '--output-dir', str(temp_dir)
        ])
        
        with pytest.raises(ValueError, match="Mask file does not exist"):
            validate_args(args)
    
    def test_validate_invalid_workers(self, sample_image_file, temp_dir):
        """Test validation fails for invalid workers value."""
        parser = create_parser()
        args = parser.parse_args([
            str(sample_image_file),
            '--output-dir', str(temp_dir),
            '--workers', '0'
        ])
        
        with pytest.raises(ValueError, match="--workers must be positive"):
            validate_args(args)
    
    def test_validate_invalid_spacing(self, sample_image_file, temp_dir):
        """Test validation fails for invalid spacing values."""
        parser = create_parser()
        args = parser.parse_args([
            str(sample_image_file),
            '--output-dir', str(temp_dir),
            '--spacing', '1.0'  # Only 1 value
        ])
        
        with pytest.raises(ValueError, match="--spacing must have 2, 3, or 4 values"):
            validate_args(args)
    
    def test_validate_conflicting_verbose_quiet(self, sample_image_file, temp_dir):
        """Test validation fails when both verbose and quiet are set."""
        parser = create_parser()
        args = parser.parse_args([
            str(sample_image_file),
            '--output-dir', str(temp_dir),
            '--verbose',
            '--quiet'
        ])
        
        with pytest.raises(ValueError, match="Cannot use both --verbose and --quiet"):
            validate_args(args)
    
    def test_validate_mask_in_batch_mode(self, sample_csv_file, temp_dir):
        """Test validation fails when --mask is used with CSV input."""
        parser = create_parser()
        args = parser.parse_args([
            str(sample_csv_file),
            '--mask', 'mask.nii.gz',
            '--output-dir', str(temp_dir)
        ])
        
        with pytest.raises(ValueError, match="only valid for single file mode"):
            validate_args(args)
    
    def test_validate_creates_output_dir(self, sample_image_file, temp_dir):
        """Test that validation creates output directory if it doesn't exist."""
        output_dir = temp_dir / "new_output"
        assert not output_dir.exists()
        
        parser = create_parser()
        args = parser.parse_args([
            str(sample_image_file),
            '--output-dir', str(output_dir)
        ])
        
        validate_args(args)
        assert output_dir.exists()


class TestModeDetection:
    """Test mode detection logic."""
    
    def test_detect_single_mode(self):
        """Test detection of single file mode."""
        assert detect_mode('image.nii.gz') == 'single'
        assert detect_mode('image.nii') == 'single'
        assert detect_mode('image.png') == 'single'
    
    def test_detect_batch_mode(self):
        """Test detection of batch mode."""
        assert detect_mode('cases.csv') == 'batch'
        assert detect_mode('data.CSV') == 'batch'


class TestConfigSaving:
    """Test configuration saving."""
    
    def test_save_config(self, temp_dir):
        """Test saving configuration to YAML file."""
        parser = create_parser()
        args = parser.parse_args([
            'image.nii.gz',
            '--output-dir', str(temp_dir),
            '--normalize',
            '--spacing', '1.0', '1.0', '2.0',
            '--methods', 'persistence_stats', 'betti_curve'
        ])
        
        config_path = temp_dir / "saved_config.yaml"
        save_config(args, config_path)
        
        assert config_path.exists()
        
        # Load and verify
        with open(config_path, 'r') as f:
            saved_config = yaml.safe_load(f)
        
        assert saved_config['preprocessing']['normalize'] is True
        assert saved_config['preprocessing']['spacing'] == [1.0, 1.0, 2.0]
        assert saved_config['vectorization']['methods'] == ['persistence_stats', 'betti_curve']


class TestMainFunction:
    """Test main CLI entry point with mocked execution."""
    
    def test_main_help(self):
        """Test that --help works."""
        with patch('sys.argv', ['medtda', '--help']):
            with pytest.raises(SystemExit) as exc_info:
                main()
            assert exc_info.value.code == 0
    
    def test_main_version(self):
        """Test that --version works."""
        with patch('sys.argv', ['medtda', '--version']):
            with pytest.raises(SystemExit) as exc_info:
                main()
            assert exc_info.value.code == 0
    
    def test_main_missing_args(self):
        """Test main function with missing required arguments."""
        with patch('sys.argv', ['medtda']):
            with pytest.raises(SystemExit) as exc_info:
                main()
            assert exc_info.value.code != 0
    
    @patch('medtda.cli.process_single_file')
    def test_main_single_file_mode(self, mock_process, sample_image_file, temp_dir):
        """Test main routes to single file processing."""
        mock_process.return_value = 0
        
        with patch('sys.argv', [
            'medtda',
            str(sample_image_file),
            '--output-dir', str(temp_dir)
        ]):
            exit_code = main()
            assert exit_code == 0
            mock_process.assert_called_once()
    
    @patch('medtda.cli.process_batch')
    def test_main_batch_mode(self, mock_process, sample_csv_file, temp_dir):
        """Test main routes to batch processing."""
        mock_process.return_value = 0
        
        with patch('sys.argv', [
            'medtda',
            str(sample_csv_file),
            '--output-dir', str(temp_dir)
        ]):
            exit_code = main()
            assert exit_code == 0
            mock_process.assert_called_once()
    
    def test_main_keyboard_interrupt(self, sample_image_file, temp_dir):
        """Test main handles keyboard interrupt."""
        with patch('medtda.cli.validate_args', side_effect=KeyboardInterrupt()):
            with patch('sys.argv', [
                'medtda',
                str(sample_image_file),
                '--output-dir', str(temp_dir)
            ]):
                exit_code = main()
                assert exit_code == 130


class TestIntegration:
    """Integration tests with actual execution (mocked FeatureExtractor)."""
    
    @patch('medtda.cli.FeatureExtractor')
    def test_single_file_execution(self, mock_extractor, sample_image_file, temp_dir):
        """Test complete single file processing flow."""
        # Mock the extractor
        mock_instance = MagicMock()
        mock_instance.execute.return_value = {
            'feature1': 1.0,
            'feature2': 2.0
        }
        mock_extractor.return_value = mock_instance
        
        with patch('sys.argv', [
            'medtda',
            str(sample_image_file),
            '--output-dir', str(temp_dir),
            '--verbose'
        ]):
            exit_code = main()
            assert exit_code == 0
        
        # Check output files
        features_csv = temp_dir / "test_image_features.csv"
        config_yaml = temp_dir / "test_image_config.yaml"
        
        assert features_csv.exists()
        assert config_yaml.exists()
        
        # Verify CSV content
        df = pd.read_csv(features_csv)
        assert len(df) == 1
        assert 'feature1' in df.columns
        assert 'feature2' in df.columns
    
    @patch('medtda.cli.FeatureExtractor')
    def test_batch_execution(self, mock_extractor, sample_csv_file, temp_dir):
        """Test complete batch processing flow."""
        # Create dummy image files
        for i in range(1, 4):
            img_file = temp_dir / f"img{i}.nii.gz"
            img_file.touch()
        
        # Update CSV with correct paths
        df = pd.DataFrame({
            'id': ['case001', 'case002', 'case003'],
            'image_path': [
                str(temp_dir / 'img1.nii.gz'),
                str(temp_dir / 'img2.nii.gz'),
                str(temp_dir / 'img3.nii.gz')
            ],
            'mask_path': ['', '', '']
        })
        csv_path = temp_dir / "cases.csv"
        df.to_csv(csv_path, index=False)
        
        # Mock the extractor
        mock_instance = MagicMock()
        mock_instance.execute.return_value = {
            'feature1': 1.0,
            'feature2': 2.0
        }
        mock_extractor.return_value = mock_instance
        
        with patch('sys.argv', [
            'medtda',
            str(csv_path),
            '--output-dir', str(temp_dir),
            '--workers', '1'
        ]):
            exit_code = main()
            assert exit_code == 0
        
        # Check output files
        batch_features = temp_dir / "batch_features.csv"
        batch_config = temp_dir / "batch_config.yaml"
        
        assert batch_features.exists()
        assert batch_config.exists()
        
        # Verify CSV content
        result_df = pd.read_csv(batch_features)
        assert len(result_df) == 3
        assert 'id' in result_df.columns
        assert list(result_df['id']) == ['case001', 'case002', 'case003']


if __name__ == '__main__':
    pytest.main([__file__, '-v'])
