"""
Module for filtering isovar results CSV files.
"""

import time
import typer
from typing import Optional, List
from typing_extensions import Annotated
from rich.progress import Progress, SpinnerColumn, TextColumn
from ..utils.isovar_filter import filter_isovar_results

app = typer.Typer()


@app.command()
def filter_isovar(
    input: Annotated[str, typer.Option("--input", "-i", help="Path to input isovar-results.csv file.")],
    output: Annotated[str, typer.Option("--output", "-o", help="Path to output filtered CSV file.")],
    columns: Annotated[Optional[List[str]], typer.Option("--columns", help="Additional columns to extract (default columns will always be included).")] = None,
    fasta: Annotated[Optional[str], typer.Option("--fasta", "-f", help="Path to output FASTA file with trimmed predicted mutant protein sequences.")] = None,
):
    """
    Filter isovar results CSV file to extract specified columns and filter rows
    where protein_sequence_contains_mutation is TRUE.
    
    This command processes isovar results and extracts the following columns by default:
    - 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
    
    Only rows where protein_sequence_contains_mutation is TRUE are included in the output.
    
    If --fasta option is provided, a FASTA file will be generated containing the
    trimmed predicted mutant protein sequences.
    """
    
    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        transient=True,
    ) as progress:
        try:
            progress.add_task(description="Filtering isovar results...", total=None)
            start = time.time()
            
            # Default columns (always included)
            default_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",
            ]
            
            # Merge default columns with any additional columns specified
            if columns:
                all_columns = list(set(default_columns + columns))
            else:
                all_columns = default_columns
            
            rows_written = filter_isovar_results(
                input=input,
                output=output,
                columns=all_columns,
                fasta=fasta
            )
            
            end = time.time()
            elapsed = end - start
            time_cost = f"{int(elapsed // 3600)}h{int((elapsed % 3600) // 60)}m{elapsed % 60:.2f}s"
            
            progress.add_task(description=f"Successfully filtered {rows_written} rows", total=None)
            print(f"Successfully filtered {rows_written} rows to {output}")
            if fasta:
                print(f"FASTA file with trimmed mutant protein sequences written to {fasta}")
            print(f"Time cost: {time_cost}")
            
        except Exception as e:
            print(f"Error: {e}")
            progress.add_task(description="Filtering failed", total=None)
            raise typer.Exit(code=1)
