"""
FASTA utilities with multiple backends
"""
import gzip
import logging
from pathlib import Path
from typing import Dict, Generator, List, Optional, Union

from .polars_utils import (
    POLARS_AVAILABLE, POLARS_BIO_AVAILABLE, BIOPYTHON_AVAILABLE,
    pl_read_fasta
)

LOG = logging.getLogger(__name__)


class FastaRecord:
    """Unified FASTA record representation."""
    
    def __init__(
        self, 
        seq_id: str, 
        sequence: str, 
        description: str = "", 
        metadata: Optional[Dict] = None
    ):
        self.seq_id = seq_id
        self.sequence = sequence
        self.description = description
        self.metadata = metadata or {}
    
    @property
    def header(self) -> str:
        """Get complete header."""
        if self.description:
            return f"{self.seq_id} {self.description}"
        return self.seq_id
    
    @property
    def length(self) -> int:
        """Get sequence length."""
        return len(self.sequence)
    
    def __str__(self) -> str:
        """FASTA format string."""
        return f">{self.header}\n{self.sequence}"


def read_fasta(
    filepath: Union[str, Path], 
    use_polars_bio: bool = True,
    use_biopython: bool = True
) -> Generator[FastaRecord, None, None]:
    """
    Read FASTA file with automatic backend selection.
    
    Args:
        filepath: Path to FASTA file
        use_polars_bio: Use polars-bio if available
        use_biopython: Use BioPython if available
    
    Yields:
        FastaRecord objects
    """
    filepath = Path(filepath)
    
    # Try polars-bio first (fastest)
    if use_polars_bio and POLARS_BIO_AVAILABLE and POLARS_AVAILABLE and pl_read_fasta:
        try:
            import polars as pl
            
            LOG.debug(f"Reading {filepath} with polars-bio")
            # Convert Path to string for polars-bio
            df = pl_read_fasta(str(filepath))
            
            for row in df.iter_rows(named=True):
                seq_id = row.get('id', '')
                description = row.get('description', '')
                sequence = row.get('sequence', '')
                
                metadata = {k: v for k, v in row.items() 
                           if k not in ['id', 'description', 'sequence']}
                
                yield FastaRecord(seq_id, sequence, description, metadata)
            
            return
        except Exception as e:
            LOG.warning(f"polars-bio failed: {e}. Falling back.")
    
    # Try BioPython
    if use_biopython and BIOPYTHON_AVAILABLE:
        try:
            from Bio import SeqIO
            
            LOG.debug(f"Reading {filepath} with BioPython")
            
            # Handle compression
            if filepath.suffix == '.gz':
                handle = gzip.open(filepath, 'rt')
            else:
                handle = open(filepath, 'r')
            
            with handle:
                for record in SeqIO.parse(handle, "fasta"):
                    yield FastaRecord(
                        seq_id=record.id,
                        sequence=str(record.seq),
                        description=record.description[len(record.id):].lstrip(),
                        metadata={
                            'name': record.name,
                            'dbxrefs': record.dbxrefs,
                            'annotations': record.annotations
                        }
                    )
            
            return
        except Exception as e:
            LOG.warning(f"BioPython failed: {e}. Falling back.")
    
    # Fallback to manual parser
    LOG.debug(f"Reading {filepath} with manual parser")
    yield from _read_fasta_manual(filepath)

def _read_fasta_manual(filepath: Path) -> Generator[FastaRecord, None, None]:
    """Manual FASTA parser for when no libraries are available."""
    # Handle compression
    if filepath.suffix == '.gz':
        handle = gzip.open(filepath, 'rt')
    elif filepath.suffix == '.bz2':
        import bz2
        handle = bz2.open(filepath, 'rt')
    else:
        handle = open(filepath, 'r')
    
    with handle:
        current_header = None
        current_sequence = []
        
        for line in handle:
            line = line.strip()
            
            if not line:
                continue
            
            if line.startswith('>'):
                # Yield previous record
                if current_header is not None and current_sequence:
                    seq_id, description = _parse_header(current_header)
                    yield FastaRecord(
                        seq_id=seq_id,
                        sequence=''.join(current_sequence),
                        description=description
                    )
                
                # Start new record
                current_header = line[1:]
                current_sequence = []
            else:
                current_sequence.append(line)
        
        # Don't forget the last record
        if current_header is not None and current_sequence:
            seq_id, description = _parse_header(current_header)
            yield FastaRecord(
                seq_id=seq_id,
                sequence=''.join(current_sequence),
                description=description
            )


def _parse_header(header: str) -> tuple:
    """Parse FASTA header into ID and description."""
    parts = header.split(None, 1)
    if len(parts) == 1:
        return parts[0], ""
    return parts[0], parts[1]


def write_fasta(
    records: List[FastaRecord], 
    filepath: Union[str, Path]
) -> None:
    """
    Write FASTA records to file.
    
    Args:
        records: List of FastaRecord objects
        filepath: Output file path
    """
    filepath = Path(filepath)
    filepath.parent.mkdir(parents=True, exist_ok=True)
    
    with open(filepath, 'w') as f:
        for record in records:
            f.write(str(record) + '\n')
    
    LOG.info(f"Written {len(records)} sequences to {filepath}")


def fasta_to_dataframe(
    filepath: Union[str, Path]
) -> Union['pl.DataFrame', List[Dict]]:
    """
    Convert FASTA file to DataFrame.
    
    Args:
        filepath: Path to FASTA file
    
    Returns:
        DataFrame with sequence information
    """
    filepath = Path(filepath)
    
    # Try polars-bio first
    if POLARS_BIO_AVAILABLE and POLARS_AVAILABLE and pl_read_fasta:
        try:
            import polars as pl
            
            LOG.debug(f"Reading {filepath} with polars-bio")
            df = pl_read_fasta(str(filepath))
            
            # Add length column
            df = df.with_columns(
                pl.col('sequence').str.lengths().alias('length')
            )
            
            return df
        except Exception as e:
            LOG.warning(f"polars-bio failed: {e}. Falling back.")
    
    # Fallback to iterative reading
    records = []
    
    for record in read_fasta(filepath):
        records.append({
            'seq_id': record.seq_id,
            'description': record.description,
            'sequence': record.sequence,
            'length': record.length,
            **record.metadata
        })
    
    # Create DataFrame if Polars is available
    if POLARS_AVAILABLE:
        import polars as pl
        
        schema = {
            'seq_id': pl.Utf8,
            'description': pl.Utf8,
            'sequence': pl.Utf8,
            'length': pl.Int64
        }
        
        return pl.DataFrame(records, schema=schema)
    
    return records


def dataframe_to_fasta(
    df: Union['pl.DataFrame', List[Dict]], 
    filepath: Union[str, Path],
    id_col: str = 'seq_id',
    desc_col: Optional[str] = 'description',
    seq_col: str = 'sequence'
) -> None:
    """
    Convert DataFrame to FASTA file.
    
    Args:
        df: DataFrame or list of dicts
        filepath: Output file path
        id_col: Column containing sequence IDs
        desc_col: Column containing descriptions (optional)
        seq_col: Column containing sequences
    """
    filepath = Path(filepath)
    filepath.parent.mkdir(parents=True, exist_ok=True)
    
    with open(filepath, 'w') as f:
        if POLARS_AVAILABLE:
            import polars as pl
            if isinstance(df, pl.DataFrame):
                for row in df.iter_rows(named=True):
                    seq_id = row.get(id_col, '')
                    description = row.get(desc_col, '') if desc_col else ''
                    sequence = row.get(seq_col, '')
                    
                    if description:
                        f.write(f">{seq_id} {description}\n")
                    else:
                        f.write(f">{seq_id}\n")
                    
                    f.write(f"{sequence}\n")
                return
        
        # Fallback for lists
        for row in df:
            if isinstance(row, dict):
                seq_id = row.get(id_col, '')
                description = row.get(desc_col, '') if desc_col else ''
                sequence = row.get(seq_col, '')
            else:
                # Assume it's a tuple/list
                seq_id = row[0] if len(row) > 0 else ''
                description = row[1] if len(row) > 1 else ''
                sequence = row[2] if len(row) > 2 else ''
            
            if description:
                f.write(f">{seq_id} {description}\n")
            else:
                f.write(f">{seq_id}\n")
            
            f.write(f"{sequence}\n")
    
    LOG.info(f"Written sequences to {filepath}")


__all__ = [
    'FastaRecord',
    'read_fasta',
    'write_fasta',
    'fasta_to_dataframe',
    'dataframe_to_fasta',
    '_parse_header',
    '_read_fasta_manual'
]
