"""
Test FASTA utilities
"""
import pytest
from pathlib import Path
import gzip

from juicekeggers.utils.fasta import (
    FastaRecord, read_fasta, write_fasta,
    fasta_to_dataframe, dataframe_to_fasta
)


def test_fasta_record():
    """Test FastaRecord class."""
    record = FastaRecord(
        seq_id="test_id",
        sequence="ATCG",
        description="test description",
        metadata={"key": "value"}
    )
    
    assert record.seq_id == "test_id"
    assert record.sequence == "ATCG"
    assert record.description == "test description"
    assert record.length == 4
    assert record.metadata["key"] == "value"
    assert str(record) == ">test_id test description\nATCG"


def test_fasta_record_no_description():
    """Test FastaRecord without description."""
    record = FastaRecord(seq_id="test", sequence="ATCG")
    assert record.header == "test"
    assert str(record) == ">test\nATCG"


def test_read_fasta(sample_fasta_file):
    """Test reading FASTA file."""
    records = list(read_fasta(sample_fasta_file))
    
    assert len(records) == 3
    assert records[0].seq_id == "seq1"
    assert records[0].description == "description1"
    assert records[0].sequence == "ATGCTAGCTAGCTAGCTAGCTAGC"
    assert records[0].length == 24
    
    assert records[2].seq_id == "seq3"
    assert records[2].description == ""
    assert records[2].sequence == "TTTTTTTTTTTTTTTTTTTTTTTT"


def test_write_fasta(temp_dir):
    """Test writing FASTA file."""
    records = [
        FastaRecord("id1", "ATCG", "desc1"),
        FastaRecord("id2", "GCTA", "desc2"),
    ]
    
    output_file = temp_dir / "output.fasta"
    write_fasta(records, output_file)
    
    assert output_file.exists()
    
    # Read it back
    read_records = list(read_fasta(output_file))
    assert len(read_records) == 2
    assert read_records[0].seq_id == "id1"
    assert read_records[1].seq_id == "id2"


def test_fasta_to_dataframe(sample_fasta_file):
    """Test converting FASTA to DataFrame."""
    result = fasta_to_dataframe(sample_fasta_file)
    
    # Check if Polars is available
    try:
        import polars as pl
        if isinstance(result, pl.DataFrame):
            assert result.height == 3
            assert "seq_id" in result.columns
            assert "sequence" in result.columns
            assert "length" in result.columns
    except ImportError:
        # Fallback to list
        assert isinstance(result, list)
        assert len(result) == 3
        assert all("seq_id" in r for r in result)


def test_dataframe_to_fasta(temp_dir):
    """Test converting DataFrame to FASTA."""
    # Create test data
    test_data = [
        {"seq_id": "gene1", "description": "protein 1", "sequence": "ATCGATCG"},
        {"seq_id": "gene2", "description": "protein 2", "sequence": "CGATCGAT"},
    ]
    
    output_file = temp_dir / "from_df.fasta"
    dataframe_to_fasta(test_data, output_file)
    
    assert output_file.exists()
    
    # Read and verify
    with open(output_file, 'r') as f:
        content = f.read()
        assert ">gene1 protein 1" in content
        assert "ATCGATCG" in content
        assert ">gene2 protein 2" in content
        assert "CGATCGAT" in content


def test_read_gzipped_fasta(temp_dir):
    """Test reading gzipped FASTA."""
    # Create gzipped FASTA
    fasta_content = """>seq1
ATCG
>seq2
GCTA
"""
    
    gz_file = temp_dir / "test.fasta.gz"
    with gzip.open(gz_file, 'wt') as f:
        f.write(fasta_content)
    
    # Read it
    records = list(read_fasta(gz_file))
    assert len(records) == 2
    assert records[0].sequence == "ATCG"
    assert records[1].sequence == "GCTA"
	