"""
Utility functions for filtering and extracting data from isovar results CSV files.
"""

import pandas as pd
from typing import List, Optional


def filter_isovar_results(
    input: str,
    output: str,
    columns: Optional[List[str]] = None,
    fasta: Optional[str] = None
) -> int:
    """
    Extract specified columns from isovar results CSV and filter rows where
    protein_sequence_contains_mutation is TRUE.
    
    Args:
        input: Path to input isovar-results.csv file
        output: Path to output filtered CSV file
        columns: List of column names to extract. If None, uses default columns.
        fasta: Path to output FASTA file with trimmed predicted mutant protein sequences.
    
    Returns:
        Number of rows in the filtered output
    
    Raises:
        FileNotFoundError: If input CSV file doesn't exist
        ValueError: If required columns are missing
    """
    # Default columns to extract
    if columns is None:
        columns = [
            "variant",
            "num_reads_supporting_top_protein_sequence",
            "num_total_reads",
            "predicted_effect_gene_name",
            "predicted_effect_gene_id",
            "predicted_effect_transcript_id",
            "predicted_effect_transcript_name",
            "predicted_effect_mutant_protein_sequence",
            "trimmed_predicted_mutant_protein_sequence",
            "trimmed_reference_protein_sequence",
            "protein_sequence_contains_mutation",
        ]
    
    # Read the CSV file
    try:
        df = pd.read_csv(input)
    except FileNotFoundError:
        raise FileNotFoundError(f"Input file not found: {input}")
    
    # Check if required columns exist
    missing_columns = [col for col in columns if col not in df.columns]
    if missing_columns:
        raise ValueError(
            f"Missing required columns in CSV file: {', '.join(missing_columns)}"
        )
    
    # Extract specified columns
    df_filtered = df[columns].copy()
    
    # Filter rows where protein_sequence_contains_mutation is TRUE
    # Handle both boolean and string representations
    if "protein_sequence_contains_mutation" in df_filtered.columns:
        # Convert to boolean if needed
        if df_filtered["protein_sequence_contains_mutation"].dtype == 'object':
            # Handle string values like "True", "TRUE", "true", etc.
            # Also handle NaN/empty values as False
            df_filtered["protein_sequence_contains_mutation"] = (
                df_filtered["protein_sequence_contains_mutation"]
                .astype(str)
                .str.upper()
                .isin(["TRUE", "1", "YES", "T"])
            )
        elif df_filtered["protein_sequence_contains_mutation"].dtype == 'bool':
            # Already boolean, no conversion needed
            pass
        else:
            # Try to convert to boolean
            df_filtered["protein_sequence_contains_mutation"] = (
                df_filtered["protein_sequence_contains_mutation"].astype(bool)
            )
        
        # Filter for TRUE values
        df_filtered = df_filtered[
            df_filtered["protein_sequence_contains_mutation"] == True
        ]
    
    # Remove the filter column from output (optional, but usually we keep it)
    # If you want to remove it, uncomment the following line:
    # df_filtered = df_filtered.drop(columns=["protein_sequence_contains_mutation"])
    
    # Write to output CSV
    df_filtered.to_csv(output, index=False)
    
    # Generate FASTA file if requested
    if fasta:
        write_fasta_file(df_filtered, fasta)
    
    return len(df_filtered)


def write_fasta_file(df_filtered: pd.DataFrame, fasta: str) -> None:
    """
    Write trimmed predicted mutant protein sequences to a FASTA file.
    
    Args:
        df_filtered: Filtered DataFrame containing the isovar results
        fasta: Path to output FASTA file
    """
    with open(fasta, 'w') as f:
        for idx, row in df_filtered.iterrows():
            # Get the trimmed predicted mutant protein sequence
            sequence = row.get('trimmed_predicted_mutant_protein_sequence', '')
            
            # Skip if sequence is empty or NaN
            if pd.isna(sequence) or not sequence:
                continue
            
            # Create a descriptive header
            variant = str(row.get('variant', f'variant_{idx}')).replace(':', '_').replace('>', '_gt_')
            gene_name = str(row.get('predicted_effect_gene_name', 'unknown_gene'))
            transcript_id = str(row.get('predicted_effect_transcript_id', 'unknown_transcript'))
            
            header = f">{variant}|{gene_name}|{transcript_id}|index_{idx}"
            
            # Write to FASTA format
            f.write(f"{header}\n")
            f.write(f"{sequence}\n")
