import pandas as pd
from io import StringIO
from typing import TextIO, Optional


def seq_df_to_fasta_handle(
    df: pd.DataFrame,
    header_col: Optional[str] = None,
    description_col: Optional[str] = None,
    seq_col: str = "seq",
) -> TextIO:
    """
    Convert a pandas sequence DataFrame to a FASTA-formatted StringIO handle.

    Parameters
    ----------
    `df`: `pandas.DataFrame`, required
    `identifier_col`: `str`, optional
    `description_col`: `str`, optional
    `seq_col Column`: `str`, required

    :param df: `pandas.DataFrame`, Sequence DataFrame.
    :param header_col: `str`, Name of column to use as the fasta sequence header if not index/None. By default, the index is the header. All spaces are replaced with underscores.
    :param description_col: `str`, Name of column to use for sequence description, if 'description_col' is not None. By default, no description is used.
    :param seq_col: `str`, Column for the sequence data.

    Returns
    -------
    `StringIO`

    :return: A `StringIO` handle for the corresponding FASTA-formatted file.
    """
    fasta_io = StringIO()

    header_col = header_col.strip() if header_col else None
    description_col = description_col.strip() if description_col else None

    # Format fasta: >header desc\nseq\n
    if header_col is None:
        header = df.index.map(str).str.strip().str.replace(r"\s+", "_", regex=True)
    else:
        header = df[header_col].astype(str).str.strip().replace(r"\s+", "_", regex=True)
    if description_col is None:
        description = ""
    else:
        description = " " + df[description_col].astype(str).str.strip()

    seq = df[seq_col.strip()].astype(str).map(lambda x: x.upper()).str.strip()

    fasta_data = ">" + header + description + "\n" + seq + "\n"

    # Combine all rows into a single string and write to buffer
    fasta_io.write("".join(fasta_data))

    # Reset pointer to the beginning for the caller to read
    fasta_io.seek(0)

    return fasta_io


if __name__ == "__main__":
    # Quick example usage
    data = {"scl": ["Cytoplasm", "Nucleus"], "seq": ["MAGA", "MTYPR"]}
    test_df = pd.DataFrame(data, index=["prot_1", "prot_2"])

    handle = seq_df_to_fasta_handle(test_df)
    print(handle.read())
