#!/usr/bin/env python3
"""
Integration tests for teloclip extend command using synthetic datasets.

This module provides comprehensive integration testing for the teloclip extend
functionality using the synthetic contigs and alignments generated by the
companion scripts. Tests cover various edge cases and scenarios.

Test Categories:
1. Basic functionality tests
2. Overhang detection and extension
3. High coverage scenarios
4. One-sided overhang handling
5. Homopolymer sequence processing
6. Complex alignment patterns
7. Control scenarios (no overhangs)
8. CLI option validation
9. Error handling and edge cases

Each test uses the pre-generated synthetic datasets to ensure consistent,
reproducible testing without requiring external dependencies.
"""

from pathlib import Path
import subprocess
import tempfile
from typing import Dict, List, Optional

import pytest

# Test data paths
TEST_DATA_DIR = Path(__file__).parent / 'test_data'
SYNTHETIC_FASTA = TEST_DATA_DIR / 'synthetic_contigs.fasta'
SYNTHETIC_BAM = TEST_DATA_DIR / 'synthetic_alignments_sorted.bam'


class TeloclipExtendRunner:
    """Helper class to run teloclip extend commands with various options."""

    def __init__(self, temp_dir: Path):
        self.temp_dir = temp_dir
        self.results = {}

    def run_extend(
        self,
        genome_fasta: Path,
        alignments_bam: Path,
        output_prefix: Optional[str] = None,
        exclude_contigs: Optional[List[str]] = None,
        exclude_contigs_file: Optional[Path] = None,
        screen_terminal_bases: Optional[int] = None,
        additional_args: Optional[List[str]] = None,
    ) -> Dict:
        """
        Run teloclip extend command with specified parameters.

        Returns
        -------
        Dict
            Command result with returncode, stdout, stderr, and output files
        """
        # Build command (BAM file first, then FASTA)
        cmd = ['teloclip', 'extend', str(alignments_bam), str(genome_fasta)]

        # Always specify an output file to ensure predictable behavior
        prefix = output_prefix or 'teloclip_extend'
        output_file = self.temp_dir / f'{prefix}_extended.fasta'
        cmd.extend(['--output-fasta', str(output_file)])

        # Add optional parameters
        if output_prefix:
            cmd.extend(['--prefix', output_prefix])

        if exclude_contigs:
            cmd.extend(['--exclude-contigs', ','.join(exclude_contigs)])

        if exclude_contigs_file:
            cmd.extend(['--exclude-contigs-file', str(exclude_contigs_file)])

        if screen_terminal_bases is not None:
            cmd.extend(['--screen-terminal-bases', str(screen_terminal_bases)])

        if additional_args:
            cmd.extend(additional_args)  # Run command in temp directory
        result = subprocess.run(
            cmd,
            cwd=self.temp_dir,
            capture_output=True,
            text=True,
            timeout=60,  # Prevent hanging tests
        )

        # Find output files
        output_files = {
            'extended_fasta': None,
            'polished_fasta': None,
            'log_file': None,
        }

        # Since we always specify --output-fasta, the output file should be predictable
        if output_file.exists():
            output_files['extended_fasta'] = output_file

        # Look for other files with the prefix
        prefix = output_prefix or 'teloclip_extend'
        for file_path in self.temp_dir.glob(f'{prefix}*'):
            if file_path.suffix == '.fasta' and 'polished' in file_path.name:
                output_files['polished_fasta'] = file_path
            elif file_path.suffix == '.log':
                output_files['log_file'] = file_path

        return {
            'returncode': result.returncode,
            'stdout': result.stdout,
            'stderr': result.stderr,
            'output_files': output_files,
            'command': ' '.join(cmd),
        }


def read_fasta_sequences(fasta_path: Path) -> Dict[str, str]:
    """Read FASTA file and return dictionary of sequences."""
    sequences = {}
    current_id = None
    current_seq = []

    with open(fasta_path, 'r') as f:
        for line in f:
            line = line.strip()
            if line.startswith('>'):
                if current_id is not None:
                    sequences[current_id] = ''.join(current_seq)
                current_id = line[1:].split()[0]  # Take first part of header
                current_seq = []
            else:
                current_seq.append(line)

    # Add last sequence
    if current_id is not None:
        sequences[current_id] = ''.join(current_seq)

    return sequences


def count_telomeric_motifs(sequence: str, motif: str = 'TTAGGG') -> int:
    """Count occurrences of telomeric motif in sequence."""
    count = 0
    for i in range(len(sequence) - len(motif) + 1):
        if sequence[i : i + len(motif)] == motif:
            count += 1
    return count


@pytest.fixture
def temp_dir():
    """Create temporary directory for test outputs."""
    with tempfile.TemporaryDirectory() as tmpdir:
        yield Path(tmpdir)


@pytest.fixture
def runner(temp_dir):
    """Create TeloclipExtendRunner instance."""
    return TeloclipExtendRunner(temp_dir)


@pytest.fixture(scope='session')
def test_data_available():
    """Check if test data files are available."""
    if not TEST_DATA_DIR.exists():
        pytest.skip('Test data directory not found')

    if not SYNTHETIC_FASTA.exists():
        pytest.skip('Synthetic FASTA file not found - run test_generate_data.py first')

    if not SYNTHETIC_BAM.exists():
        pytest.skip(
            'Synthetic BAM file not found - run test_generate_alignments.py first'
        )

    return True


class TestBasicFunctionality:
    """Test basic teloclip extend functionality."""

    def test_basic_extend_runs_successfully(self, runner, test_data_available):
        """Test that basic extend command runs without errors."""
        result = runner.run_extend(SYNTHETIC_FASTA, SYNTHETIC_BAM)

        assert result['returncode'] == 0, f'Command failed: {result["stderr"]}'
        assert result['output_files']['extended_fasta'] is not None
        assert result['output_files']['extended_fasta'].exists()

    def test_custom_prefix_option(self, runner, test_data_available):
        """Test that --prefix option works correctly."""
        custom_prefix = 'test_custom_prefix'
        result = runner.run_extend(
            SYNTHETIC_FASTA, SYNTHETIC_BAM, output_prefix=custom_prefix
        )

        assert result['returncode'] == 0
        # Check that output files use custom prefix
        assert any(
            custom_prefix in str(f) for f in result['output_files'].values() if f
        )

    def test_output_file_contains_all_contigs(self, runner, test_data_available):
        """Test that output FASTA contains all expected contigs."""
        result = runner.run_extend(SYNTHETIC_FASTA, SYNTHETIC_BAM)

        assert result['returncode'] == 0

        # Read original and extended sequences
        original_seqs = read_fasta_sequences(SYNTHETIC_FASTA)
        extended_seqs = read_fasta_sequences(result['output_files']['extended_fasta'])

        # All original contigs should be in output
        assert set(original_seqs.keys()).issubset(set(extended_seqs.keys()))

        # Extended sequences should be same length or longer
        for contig_id in original_seqs:
            if contig_id in extended_seqs:
                assert len(extended_seqs[contig_id]) >= len(original_seqs[contig_id])


class TestOverhangDetection:
    """Test overhang detection and extension functionality."""

    def test_contig1_normal_overhangs(self, runner, test_data_available):
        """Test normal overhang detection on contig_1."""
        result = runner.run_extend(SYNTHETIC_FASTA, SYNTHETIC_BAM)

        assert result['returncode'] == 0

        original_seqs = read_fasta_sequences(SYNTHETIC_FASTA)
        extended_seqs = read_fasta_sequences(result['output_files']['extended_fasta'])

        # contig_1 should have extensions due to overhang alignments
        original_len = len(original_seqs['contig_1'])
        extended_len = len(extended_seqs['contig_1'])

        assert extended_len > original_len, (
            f'contig_1 should be extended: {original_len} -> {extended_len}'
        )

    def test_contig2_high_coverage_extensions(self, runner, test_data_available):
        """Test high coverage overhang handling on contig_2."""
        result = runner.run_extend(SYNTHETIC_FASTA, SYNTHETIC_BAM)

        assert result['returncode'] == 0

        original_seqs = read_fasta_sequences(SYNTHETIC_FASTA)
        extended_seqs = read_fasta_sequences(result['output_files']['extended_fasta'])

        # contig_2 has 10x coverage, should have significant extensions
        original_len = len(original_seqs['contig_2'])
        extended_len = len(extended_seqs['contig_2'])

        assert extended_len > original_len, (
            f'contig_2 should be extended: {original_len} -> {extended_len}'
        )

        # Should have more extension than contig_1 due to higher coverage
        contig1_extension = len(extended_seqs['contig_1']) - len(
            original_seqs['contig_1']
        )
        contig2_extension = extended_len - original_len

        # Note: This assertion may not always be true depending on consensus algorithm
        # But it's a reasonable expectation for this test data
        print(
            f'contig_1 extension: {contig1_extension}, contig_2 extension: {contig2_extension}'
        )

    def test_contig3_one_sided_overhangs(self, runner, test_data_available):
        """Test one-sided overhang handling on contig_3."""
        result = runner.run_extend(SYNTHETIC_FASTA, SYNTHETIC_BAM)

        assert result['returncode'] == 0

        original_seqs = read_fasta_sequences(SYNTHETIC_FASTA)
        extended_seqs = read_fasta_sequences(result['output_files']['extended_fasta'])

        # contig_3 should be extended (left side only has overhangs)
        original_len = len(original_seqs['contig_3'])
        extended_len = len(extended_seqs['contig_3'])

        assert extended_len > original_len, (
            f'contig_3 should be extended: {original_len} -> {extended_len}'
        )


class TestMotifProcessing:
    """Test telomeric motif processing and screening."""

    def test_contig5_motif_rich_extensions(self, runner, test_data_available):
        """Test motif-rich contig processing."""
        result = runner.run_extend(SYNTHETIC_FASTA, SYNTHETIC_BAM)

        assert result['returncode'] == 0

        extended_seqs = read_fasta_sequences(result['output_files']['extended_fasta'])

        # contig_5 is motif-rich, check for telomeric motifs in extensions
        contig5_seq = extended_seqs['contig_5']
        ttaggg_count = count_telomeric_motifs(contig5_seq, 'TTAGGG')
        ccctaa_count = count_telomeric_motifs(contig5_seq, 'CCCTAA')

        # Should find telomeric motifs (original + extensions)
        assert ttaggg_count > 0 or ccctaa_count > 0, (
            'Should find telomeric motifs in contig_5'
        )

    def test_terminal_base_screening(self, runner, test_data_available):
        """Test --screen-terminal-bases option."""
        # Test with terminal base screening enabled
        result = runner.run_extend(
            SYNTHETIC_FASTA, SYNTHETIC_BAM, screen_terminal_bases=10
        )

        assert result['returncode'] == 0

        # Should produce output (specific behavior depends on implementation)
        assert result['output_files']['extended_fasta'].exists()


class TestContigExclusion:
    """Test contig exclusion functionality."""

    def test_exclude_contigs_option(self, runner, temp_dir, test_data_available):
        """Test --exclude-contigs option."""
        excluded_contigs = ['contig_6', 'contig_7']
        result = runner.run_extend(
            SYNTHETIC_FASTA, SYNTHETIC_BAM, exclude_contigs=excluded_contigs
        )

        assert result['returncode'] == 0

        # Read original and extended sequences
        original_seqs = read_fasta_sequences(SYNTHETIC_FASTA)
        extended_seqs = read_fasta_sequences(result['output_files']['extended_fasta'])

        # All contigs (including excluded ones) should be in output
        for contig_name in original_seqs.keys():
            assert contig_name in extended_seqs, (
                f'Contig {contig_name} missing from output'
            )

        # Excluded contigs should NOT be extended (same length as original)
        for excluded in excluded_contigs:
            original_len = len(original_seqs[excluded])
            extended_len = len(extended_seqs[excluded])
            assert extended_len == original_len, (
                f'Excluded contig {excluded} was extended: {original_len} -> {extended_len}'
            )

        # Non-excluded contigs should potentially be extended (length >= original)
        non_excluded = ['contig_1', 'contig_2', 'contig_3', 'contig_4', 'contig_5']
        for included in non_excluded:
            original_len = len(original_seqs[included])
            extended_len = len(extended_seqs[included])
            assert extended_len >= original_len, (
                f'Non-excluded contig {included} became shorter: {original_len} -> {extended_len}'
            )

    def test_exclude_contigs_file(self, runner, temp_dir, test_data_available):
        """Test --exclude-contigs-file option."""
        # Create exclusion file
        exclude_file = temp_dir / 'exclude.txt'
        with open(exclude_file, 'w') as f:
            f.write('contig_6\ncontig_7\n')

        result = runner.run_extend(
            SYNTHETIC_FASTA, SYNTHETIC_BAM, exclude_contigs_file=exclude_file
        )

        assert result['returncode'] == 0

        # Read original and extended sequences
        original_seqs = read_fasta_sequences(SYNTHETIC_FASTA)
        extended_seqs = read_fasta_sequences(result['output_files']['extended_fasta'])

        # All contigs (including excluded ones) should be in output
        for contig_name in original_seqs.keys():
            assert contig_name in extended_seqs, (
                f'Contig {contig_name} missing from output'
            )

        # Excluded contigs should NOT be extended (same length as original)
        excluded_contigs = ['contig_6', 'contig_7']
        for excluded in excluded_contigs:
            original_len = len(original_seqs[excluded])
            extended_len = len(extended_seqs[excluded])
            assert extended_len == original_len, (
                f'Excluded contig {excluded} was extended: {original_len} -> {extended_len}'
            )


class TestControlScenarios:
    """Test control scenarios and edge cases."""

    def test_contig6_no_extensions(self, runner, test_data_available):
        """Test contig_6 (control) has no or minimal extensions."""
        result = runner.run_extend(SYNTHETIC_FASTA, SYNTHETIC_BAM)

        assert result['returncode'] == 0

        original_seqs = read_fasta_sequences(SYNTHETIC_FASTA)
        extended_seqs = read_fasta_sequences(result['output_files']['extended_fasta'])

        # contig_6 is control with no overhangs, should have minimal extension
        original_len = len(original_seqs['contig_6'])
        extended_len = len(extended_seqs['contig_6'])

        # Allow for small extensions (consensus might add/remove a few bases)
        extension = extended_len - original_len
        assert abs(extension) <= 10, (
            f'contig_6 should have minimal extension, got {extension}'
        )

    def test_complex_alignments_handling(self, runner, test_data_available):
        """Test handling of complex CIGAR alignments on contig_7."""
        result = runner.run_extend(SYNTHETIC_FASTA, SYNTHETIC_BAM)

        assert result['returncode'] == 0

        # Should handle complex alignments without crashing
        extended_seqs = read_fasta_sequences(result['output_files']['extended_fasta'])
        assert 'contig_7' in extended_seqs

        # Sequence should be valid DNA
        contig7_seq = extended_seqs['contig_7']
        assert all(base in 'ATGCNatgcn' for base in contig7_seq)


class TestErrorHandling:
    """Test error handling and edge cases."""

    def test_nonexistent_fasta_file(self, runner, temp_dir):
        """Test handling of non-existent FASTA file."""
        fake_fasta = temp_dir / 'nonexistent.fasta'

        result = runner.run_extend(fake_fasta, SYNTHETIC_BAM)

        # Should fail gracefully
        assert result['returncode'] != 0
        assert (
            'error' in result['stderr'].lower()
            or 'not found' in result['stderr'].lower()
        )

    def test_nonexistent_bam_file(self, runner, temp_dir):
        """Test handling of non-existent BAM file."""
        fake_bam = temp_dir / 'nonexistent.bam'

        result = runner.run_extend(SYNTHETIC_FASTA, fake_bam)

        # Should fail gracefully
        assert result['returncode'] != 0
        assert (
            'error' in result['stderr'].lower()
            or 'not found' in result['stderr'].lower()
        )

    def test_invalid_exclude_contigs_file(self, runner, temp_dir, test_data_available):
        """Test handling of invalid exclude contigs file."""
        fake_exclude_file = temp_dir / 'nonexistent_exclude.txt'

        result = runner.run_extend(
            SYNTHETIC_FASTA, SYNTHETIC_BAM, exclude_contigs_file=fake_exclude_file
        )

        # Should fail gracefully
        assert result['returncode'] != 0


class TestOutputValidation:
    """Test output file validation and quality."""

    def test_output_fasta_format(self, runner, test_data_available):
        """Test that output FASTA file is properly formatted."""
        result = runner.run_extend(SYNTHETIC_FASTA, SYNTHETIC_BAM)

        assert result['returncode'] == 0

        output_fasta = result['output_files']['extended_fasta']
        assert output_fasta.exists()

        # Check FASTA format
        with open(output_fasta, 'r') as f:
            lines = f.readlines()

        assert len(lines) > 0, 'Output file should not be empty'
        assert lines[0].startswith('>'), 'First line should be FASTA header'

        # Check for proper line lengths (should be wrapped)
        sequence_lines = [line for line in lines if not line.startswith('>')]
        if sequence_lines:
            # Most lines should be wrapped to reasonable length (e.g., 80 chars)
            line_lengths = [len(line.strip()) for line in sequence_lines]
            # Allow for last line of each sequence to be shorter
            normal_lines = [length for length in line_lengths if length > 0]
            if normal_lines:
                max_normal_length = max(normal_lines)
                assert max_normal_length <= 80, (
                    f'Lines should be wrapped, found length {max_normal_length}'
                )

    def test_log_file_creation(self, runner, test_data_available):
        """Test that log file is created and contains useful information."""
        result = runner.run_extend(SYNTHETIC_FASTA, SYNTHETIC_BAM)

        assert result['returncode'] == 0

        log_file = result['output_files']['log_file']
        if log_file and log_file.exists():
            with open(log_file, 'r') as f:
                log_content = f.read()

            # Log should contain some useful information
            assert len(log_content) > 0, 'Log file should not be empty'
            # Look for common log elements
            assert any(
                keyword in log_content.lower()
                for keyword in [
                    'processing',
                    'extending',
                    'contig',
                    'alignment',
                    'complete',
                ]
            )


# Integration test that runs all components together
class TestFullIntegration:
    """Full integration tests combining multiple features."""

    def test_full_pipeline_with_all_options(
        self, runner, temp_dir, test_data_available
    ):
        """Test complete pipeline with multiple options enabled."""
        # Create exclude file
        exclude_file = temp_dir / 'exclude_test.txt'
        with open(exclude_file, 'w') as f:
            f.write('contig_6\n')  # Exclude control contig

        result = runner.run_extend(
            SYNTHETIC_FASTA,
            SYNTHETIC_BAM,
            output_prefix='full_test',
            exclude_contigs_file=exclude_file,
            screen_terminal_bases=15,
        )

        assert result['returncode'] == 0, f'Full pipeline failed: {result["stderr"]}'

        # Validate output
        original_seqs = read_fasta_sequences(SYNTHETIC_FASTA)
        extended_seqs = read_fasta_sequences(result['output_files']['extended_fasta'])

        # All contigs (including excluded ones) should be in output
        for contig_name in original_seqs.keys():
            assert contig_name in extended_seqs, (
                f'Contig {contig_name} missing from output'
            )

        # Excluded contig should NOT be extended (same length as original)
        original_len = len(original_seqs['contig_6'])
        extended_len = len(extended_seqs['contig_6'])
        assert extended_len == original_len, (
            f'Excluded contig_6 was extended: {original_len} -> {extended_len}'
        )

        # Output files should use custom prefix
        assert 'full_test' in str(result['output_files']['extended_fasta'])


if __name__ == '__main__':
    # Run tests when script is executed directly
    pytest.main([__file__, '-v'])
