import csv
import gzip
from io import BufferedWriter, BytesIO
from typing import List, Optional, Sequence, Tuple, Union

import pandas as pd

from countess import VERSION
from countess.core.logger import Logger
from countess.core.parameters import (
    ArrayParam,
    BooleanParam,
    ChoiceParam,
    DataTypeOrNoneChoiceParam,
    FileSaveParam,
    MultiParam,
    StringParam,
)
from countess.core.plugins import PandasInputFilesPlugin, PandasOutputPlugin
from countess.utils.pandas import flatten_columns

CSV_FILE_TYPES: Sequence[Tuple[str, Union[str, List[str]]]] = [
    ("CSV", [".csv", ".csv.gz"]),
    ("TSV", [".tsv", ".tsv.gz"]),
    ("TXT", [".txt", ".txt.gz"]),
]


class LoadCsvPlugin(PandasInputFilesPlugin):
    """Load CSV files"""

    name = "CSV Load"
    description = "Loads data from CSV or similar delimited text files and assigns types to columns"
    link = "https://countess-project.github.io/CountESS/included-plugins/#csv-reader"
    version = VERSION
    file_types = CSV_FILE_TYPES

    parameters = {
        "delimiter": ChoiceParam("Delimiter", ",", choices=[",", ";", "TAB", "|", "WHITESPACE"]),
        "quoting": ChoiceParam("Quoting", "None", choices=["None", "Double-Quote", "Quote with Escape"]),
        "comment": ChoiceParam("Comment", "None", choices=["None", "#", ";"]),
        "header": BooleanParam("CSV file has header row?", True),
        "filename_column": StringParam("Filename Column", ""),
        "columns": ArrayParam(
            "Columns",
            MultiParam(
                "Column",
                {
                    "name": StringParam("Column Name", ""),
                    "type": DataTypeOrNoneChoiceParam("Column Type"),
                    "index": BooleanParam("Index?", False),
                },
            ),
        ),
    }

    def read_file_to_dataframe(self, file_params, logger, row_limit=None):
        filename = file_params["filename"].value

        options = {
            "header": 0 if self.parameters["header"].value else None,
        }
        if row_limit is not None:
            options["nrows"] = row_limit

        index_col_numbers = []

        if len(self.parameters["columns"]):
            options["names"] = []
            options["dtype"] = {}
            options["usecols"] = []

            for n, pp in enumerate(self.parameters["columns"]):
                options["names"].append(pp["name"].value or f"column_{n}")
                if pp["type"].is_not_none():
                    if pp["index"].value:
                        index_col_numbers.append(len(options["usecols"]))
                    options["usecols"].append(n)
                    options["dtype"][n] = pp["type"].get_selected_type()

        delimiter = self.parameters["delimiter"].value
        if delimiter == "TAB":
            options["delimiter"] = "\t"
        elif delimiter == "WHITESPACE":
            options["delim_whitespace"] = True
        else:
            options["delimiter"] = delimiter

        quoting = self.parameters["quoting"].value
        if quoting == "None":
            options["quoting"] = csv.QUOTE_NONE
        elif quoting == "Double-Quote":
            options["quotechar"] = '"'
            options["doublequote"] = True
        elif quoting == "Quote with Escape":
            options["quotechar"] = '"'
            options["doublequote"] = False
            options["escapechar"] = "\\"

        comment = self.parameters["comment"].value
        if comment != "None":
            options["comment"] = comment

        # XXX pd.read_csv(index_col=) is half the speed of pd.read_csv().set_index()

        df = pd.read_csv(filename, **options)

        while len(df.columns) > len(self.parameters["columns"]):
            self.parameters["columns"].add_row()

        if self.parameters["header"].value:
            for n, col in enumerate(df.columns):
                if not self.parameters["columns"][n]["name"].value:
                    self.parameters["columns"][n]["name"].value = str(col)
                    self.parameters["columns"][n]["type"].value = "string"

        filename_column = self.parameters["filename_column"].value
        if filename_column:
            df[filename_column] = filename

        if index_col_numbers:
            df = df.set_index([df.columns[n] for n in index_col_numbers])

        return df


class SaveCsvPlugin(PandasOutputPlugin):
    name = "CSV Save"
    description = "Save data as CSV or similar delimited text files"
    link = "https://countess-project.github.io/CountESS/included-plugins/#csv-writer"
    version = VERSION
    file_types = CSV_FILE_TYPES

    parameters = {
        "header": BooleanParam("CSV header row?", True),
        "filename": FileSaveParam("Filename", file_types=file_types),
        "delimiter": ChoiceParam("Delimiter", ",", choices=[",", ";", "TAB", "|", "SPACE"]),
        "quoting": BooleanParam("Quote all Strings", False),
    }

    filehandle: Optional[Union[BufferedWriter, BytesIO]] = None
    csv_columns = None

    SEPARATORS = {",": ",", ";": ";", "SPACE": " ", "TAB": "\t"}
    QUOTING = {False: csv.QUOTE_MINIMAL, True: csv.QUOTE_NONNUMERIC}

    def prepare(self, sources: list[str], row_limit: Optional[int] = None):
        if row_limit is None:
            filename = self.parameters["filename"].value
            if filename.endswith(".gz"):
                self.filehandle = gzip.open(filename, "wb")
            else:
                self.filehandle = open(filename, "wb")
        else:
            self.filehandle = BytesIO()

        self.csv_columns = None

    def process(self, data: pd.DataFrame, source: str, logger: Logger):
        # reset indexes so we can treat all columns equally.
        # if there's just a nameless index then we don't care about it, drop it.
        drop_index = data.index.name is None and data.index.names[0] is None
        dataframe = flatten_columns(data.reset_index(drop=drop_index))

        # if this is our first dataframe to write then decide whether to
        # include the header or not.
        if self.csv_columns is None:
            self.csv_columns = list(dataframe.columns)
            emit_header = bool(self.parameters["header"].value)
        else:
            # add in any columns we haven't seen yet in previous dataframes.
            for c in dataframe.columns:
                if c not in self.csv_columns:
                    self.csv_columns.append(c)
                    logger.warning(f"Added CSV Column {repr(c)} with no header")
            # fill in blanks for any columns which are in previous dataframes but not
            # in this one.
            dataframe = dataframe.assign(**{c: None for c in self.csv_columns if c not in dataframe.columns})
            emit_header = False

        dataframe.to_csv(
            self.filehandle,
            header=emit_header,
            columns=self.csv_columns,
            index=False,
            sep=self.SEPARATORS[self.parameters["delimiter"].value],
            quoting=self.QUOTING[self.parameters["quoting"].value],
        )  # type: ignore [call-overload]
        return []

    def finalize(self, logger: Logger):
        if isinstance(self.filehandle, BytesIO):
            yield self.filehandle.getvalue().decode("utf-8")
