# ta_ml/data/file_handler.py
"""
A utility class for handling file operations, such as parsing uploaded files.
"""
from io import BytesIO
from typing import IO, Literal, Optional

import pandas as pd
from fastapi import UploadFile


class FileHandlerError(Exception):
    """Custom exception for file handling errors."""
    pass


class FileHandler:
    """
    A class providing static methods to handle file reading and writing.
    """

    @staticmethod
    def read_to_dataframe(file: UploadFile) -> pd.DataFrame:
        """
        Reads an uploaded file (CSV or Excel) into a pandas DataFrame.

        Args:
            file (UploadFile): The file uploaded via a FastAPI endpoint.

        Returns:
            pd.DataFrame: The content of the file as a pandas DataFrame.

        Raises:
            FileHandlerError: If the file format is unsupported or if there's
                              a reading error.
        """
        if file.filename.endswith(".csv"):
            parser = pd.read_csv
        elif file.filename.endswith((".xls", ".xlsx")):
            parser = pd.read_excel
        else:
            raise FileHandlerError(
                "Unsupported file format. Please use CSV or Excel (.xls, .xlsx)."
            )

        try:
            return parser(file.file)
        except Exception as e:
            raise FileHandlerError(f"Error reading file: {e}") from e

    @staticmethod
    def export_dataframe(
        df: pd.DataFrame,
        format: Literal["csv", "xlsx"]
    ) -> tuple[BytesIO, str, str]:
        """
        Exports a pandas DataFrame to an in-memory file (CSV or XLSX).

        Args:
            df (pd.DataFrame): The DataFrame to export.
            format (Literal["csv", "xlsx"]): The desired output format.

        Returns:
            A tuple containing:
            - BytesIO: The in-memory buffer containing the file data.
            - str: The appropriate MIME type for the format.
            - str: The suggested filename for the download.
        """
        buffer = BytesIO()
        if format == "csv":
            df.to_csv(buffer, index=False)
            media_type = "text/csv"
            filename = "export.csv"
        elif format == "xlsx":
            with pd.ExcelWriter(buffer, engine="openpyxl") as writer:
                df.to_excel(writer, index=False, sheet_name="Export")
            media_type = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
            filename = "export.xlsx"
        else:
            raise ValueError("Unsupported export format specified.")

        buffer.seek(0)
        return buffer, media_type, filename