from __future__ import annotations

import contextlib
import itertools
import io
import functools
import os
from typing import Any, Collection, List, Dict, Union, Type, Iterable, TypeVar, Callable, Iterator, cast
import pathlib

import tabulate
import pandas as pd
from pandas.io.sql import SQLTable, pandasSQL_builder
from pandas.core.indexes.base import Index
import numpy as np
from maybe import Maybe

from .str import Str
from .enums import Enum
from .datetime import DateTime


pd.set_option("max_columns", None)

PathLike = Union[str, os.PathLike, pathlib.Path]
FuncSig = TypeVar("FuncSig", bound=Callable)


def _check_import_is_available(func: FuncSig) -> FuncSig:
    @functools.wraps(func)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        try:
            return func(*args, **kwargs)
        except ModuleNotFoundError as ex:
            exception = ModuleNotFoundError(f"Use of method '{func.__name__}' requires module '{ex.name}'. Please ensure it is available by installing it with pip.")
            exception.name, exception.path = ex.name, ex.path
            raise exception

    return cast(FuncSig, wrapper)


class Frame(pd.DataFrame):
    columns: Union[Index, Iterable]

    class InferRange(Enum):
        TRIM_SURROUNDING, STRIP_NULLS, SMALLEST_VALID = "trim_surrounding", "strip_nulls", "smallest_valid"

    class ColumnCase(Enum):
        IGNORE, SNAKE, CAMEL, PASCAL = "ignore", "snake", "camel", "pascal"

    class PathType(Enum):
        PATHMAGIC, PATHLIB, STRING = "pathmagic", "pathlib", "string"

    DEFAULT_COLUMN_CASE = ColumnCase.SNAKE
    DEFAULT_INFER_RANGE = InferRange.SMALLEST_VALID
    DEFAULT_PATH_TYPE = PathType.PATHMAGIC
    DEFAULT_SHEET_NAME = "Sheet1"

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        with self._using_parent_constructor():
            super().__init__(*args, **kwargs)
            self._clean_dtypes()

    def __repr__(self) -> str:
        return f"{type(self).__name__}(rows={len(self.index)}, columns={list(self.columns)})"

    @property
    def _constructor(self) -> Type[pd.DataFrame]:
        return type(self) if self._using_own_constructor() else pd.DataFrame

    def to_excel(self, filepath: PathLike, sheet_name: str = None, index: bool = False, **kwargs: Any) -> PathLike:
        """Write this Frame to an xlsx file. Returns that File."""
        with ExcelWriter(filepath) as writer:
            self._write_to_excel(writer=writer, sheet_name=sheet_name, index=index, **kwargs)
        return self._get_path_constructor()(filepath)

    def to_sql(self, engine: Any, name: str, if_exists: str = "fail", index: bool = True, index_label: str = "id", primary_key: str = "id", schema: str = None, dtype: dict = None, **kwargs: Any) -> None:
        """Override of the pandas.DataFrame.to_sql() method allowing a primary key identity field to be supplied when creating the sql table."""
        table = SQLTable(name, pandasSQL_builder(engine), frame=self, index=index, if_exists=if_exists, index_label=index_label, keys=primary_key, schema=schema, dtype=dtype, **kwargs)
        table.create()
        table.insert(None)

    @_check_import_is_available
    def to_table(self, table: str, schema: str = None, database: str = None, if_exists: str = "fail", primary_key: str = "id", sql: Any = None) -> Any:
        """Load this Frame into a SQL database table using the config defaults of the sqlhandler library. An sqlhandler.Sql object can be provided to override the connection defaults."""
        from sqlhandler import Sql
        sql = sql if sql is not None else Sql(database=database)
        return sql.frame_to_table(self, table=table, schema=schema, if_exists=Sql.IfExists(if_exists), primary_key=primary_key)

    def to_dataframe(self) -> pd.DataFrame:
        """Convert this Frame to a pandas.DataFrame."""
        return pd.DataFrame(self)

    def sanitize_colnames(self, casing: str = None) -> Frame:
        """Strip newlines from the column names and apply arbitrary casing."""
        df: Frame = self.copy()
        df.columns = [str(colname).strip().replace("\n", "") for colname in df.columns]

        casing = Maybe(casing).else_(self.DEFAULT_COLUMN_CASE)
        if casing is not None:
            if casing == self.ColumnCase.SNAKE:
                df.columns = [Str(colname).case.snake() for colname in df.columns]
            elif casing == self.ColumnCase.CAMEL:
                df.columns = [Str(colname).case.camel() for colname in df.columns]
            elif casing == self.ColumnCase.PASCAL:
                df.columns = [Str(colname).case.pascal() for colname in df.columns]
            else:
                self.ColumnCase.raise_if_not_a_member(casing)

        return df

    def to_ascii(self, index: bool = False, fancy: bool = True) -> str:
        """Convert this Frame to an ascii representation."""
        return str(tabulate.tabulate(self, headers=self.columns, tablefmt="fancy_grid" if fancy else "grid", showindex="never" if not index else "default"))

    @_check_import_is_available
    def to_desktop_as_excel(self, name: str, with_timestamp: bool = True, index: bool = False, **kwargs: Any) -> File:
        """Save this Frame to the current user's desktop as an xlsx file. Returns that File."""
        from pathmagic import Dir

        file = Dir.from_desktop().new_file(f"{name}_{DateTime.today().to_filetag() if with_timestamp else ''}", "xlsx")
        self.to_excel(file.path, index=index, **kwargs)
        return file

    # noinspection PyMethodOverriding
    def pivot(self, field_col: Union[pd.Series, str], value_col: Union[pd.Series, str]) -> Frame:
        """Pivot the content of this Frame."""
        field_name = field_col.name if isinstance(field_col, pd.Series) else field_col
        value_name = value_col.name if isinstance(value_col, pd.Series) else value_col

        pivoted: Frame = self.set_index([*[colname for colname in self.columns if colname not in {field_name, value_name}], field_name]).unstack(level=field_name).reset_index()
        pivoted.columns = [list(itertools.takewhile(lambda val: val, item))[-1] for item in pivoted.columns.to_flat_index()]
        return pivoted

    def unpivot(self, index_cols: List[Union[pd.Series, str]] = None, cols_to_unpivot: List[Union[pd.Series, str]] = None, field_name: str = "field", value_name: str = "value") -> Frame:
        """Unpivot the content of this Frame."""
        indexed = self.set_index(index_cols) if index_cols is not None else self.copy()

        if cols_to_unpivot is not None:
            indexed.drop([colname for colname in indexed.columns if colname not in cols_to_unpivot], axis=1, inplace=True)

        unpivoted = type(self)(indexed.stack().reset_index())
        final: Frame = unpivoted.rename(columns={unpivoted.columns[-2]: field_name, unpivoted.columns[-1]: value_name})
        final[field_name] = final[field_name].replace(True, 1).replace(False, 0)
        final.dropna(how="all", subset=list(final.columns[:-2]), inplace=True)

        return final

    def infer_dtypes(self) -> Frame:
        """Return a new Frame with newly inferred dtypes."""
        return type(self)(self.to_dict())

    def fillna_as_none(self) -> Frame:
        """Fill any nan values with None. Returns self."""
        df = self.copy()
        for name, col in df.iteritems():
            if col.isnull().any():
                df[name] = col.astype(object)

        return cast(Frame, df.where(df.notnull(), None))

    @_check_import_is_available
    def profile_report(self, *args: Any, style: dict = None, **kwargs: Any) -> Any:
        """Produce and return a pandas profile report"""
        import pandas_profiling

        style = Maybe(style).else_({'full_width': True})
        return pandas_profiling.ProfileReport(pd.DataFrame(self), *args, style=style, **kwargs)

    def profile_report_to(self, path: PathLike, *args: Any, **kwargs: Any) -> PathLike:
        """Produce a pandas profile report and return it as an html file"""
        file = self._get_path_constructor()(path)
        self.profile_report(*args, **kwargs).to_file(output_file=str(file))

        return file

    @classmethod
    def many_to_excel(cls, frames: Union[Collection[Frame], Dict[str, Collection[Frame]]], filepath: os.PathLike, index: bool = False, **kwargs: Any) -> PathLike:
        """Write an iterable of Frames or a mapping of string-keys and Frame-values into an xlsx file with several sheets."""

        named_frames = frames if isinstance(frames, dict) else {f"Sheet{idx + 1}": frame for idx, frame in enumerate(frames)}

        with ExcelWriter(filepath=filepath) as writer:
            for name, frame in named_frames.items():
                frame._write_to_excel(writer=writer, sheet_name=name, index=index, **kwargs)

        return cls._get_path_constructor()(filepath)

    @classmethod
    def from_excel(cls, filepath: os.PathLike, casing: str = None, skipcols: int = 0, infer_headers: bool = True, infer_range: str = None, password: str = None, **kwargs: Any) -> Frame:
        """Reads in the specified Excel spreadsheet into a Frame. Passes on arguments to the pandas read_excel function. Optionally changes casing on column names and strips out non-ascii characters."""

        filepath = os.fspath(filepath)

        if password is None:
            frame = cls(pd.read_excel(filepath, **kwargs))
        else:
            with cls._decrypted_stream_from_xlsx(path=filepath, password=password) as stream:
                frame = cls(pd.read_excel(stream, **kwargs))

        if skipcols:
            frame = frame.iloc[:, skipcols:]

        if infer_headers:
            frame._infer_column_headers()

        frame = frame.sanitize_colnames(casing=Maybe(casing).else_(cls.DEFAULT_COLUMN_CASE))

        infer_range = Maybe(infer_range).else_(cls.DEFAULT_INFER_RANGE)
        if infer_range:
            frame._infer_range(mode=infer_range)

        frame._infer_boolean_columns()
        return frame

    @classmethod
    def from_csv(cls, filepath: os.PathLike, casing: str = None, skipcols: int = 0, **kwargs: Any) -> Frame:
        """Reads in the specified csv file into a Frame. Passes on arguments to the pandas read_csv function. Optionally changes casing on column names and strips out non-ascii characters."""

        frame = cls(pd.read_csv(os.fspath(filepath), **kwargs))

        if skipcols:
            frame = frame.iloc[:, skipcols:]

        frame._infer_boolean_columns()
        return frame.sanitize_colnames(casing=casing)

    @classmethod
    def from_object(cls, obj: Any, private: bool = False) -> Frame:
        """Create a Frame from a single python object"""
        mappings = vars(obj) if private else {name: attr for name, attr in vars(obj).items() if not name.startswith("_")}
        return cls({name: [val] for name, val in mappings.items()})

    @classmethod
    def from_objects(cls, objects: Iterable[Any], private: bool = False) -> Frame:
        """Create a Frame from a homogenous list of python objects"""
        attrs = list(dict.fromkeys([name for obj in objects for name in vars(obj)]))
        valid_attrs = attrs if private else [name for name in attrs if not name.startswith("_")]
        return cls([tuple(vars(obj).get(attr) for attr in valid_attrs) for obj in objects], columns=valid_attrs)

    def _using_own_constructor(self) -> bool:
        return bool(object.__getattribute__(self, "_own_constructor_"))

    def _use_own_constructor(self, own: bool) -> None:
        object.__setattr__(self, "_own_constructor_", own)

    @contextlib.contextmanager
    def _using_parent_constructor(self) -> Iterator[None]:
        self._use_own_constructor(False)
        yield
        self._use_own_constructor(True)

    @classmethod
    @_check_import_is_available
    def _get_path_constructor(cls) -> Callable[..., PathLike]:
        if cls.DEFAULT_PATH_TYPE == Frame.PathType.PATHMAGIC:
            from pathmagic import File
            return cast(Callable[..., File], File.from_pathlike)
        elif cls.DEFAULT_PATH_TYPE == Frame.PathType.PATHLIB:
            return pathlib.Path
        elif cls.DEFAULT_PATH_TYPE == Frame.PathType.STRING:
            return os.fspath
        else:
            raise cls.PathType.raise_if_not_a_member(cls.DEFAULT_PATH_TYPE)

    def _infer_column_headers(self) -> None:
        col_run = len(max("".join(["0" if self._value_is_null(col) else "1" for col in self.columns]).split("0")))

        longest_runs: Dict[int, Any] = {}
        for index, row in self.iterrows():
            longest_runs.setdefault(len(max("".join(["0" if self._value_is_null(col) else "1" for col in row]).split("0"))), []).append(index)

        if max(longest_runs) > col_run:
            first_longest_index = longest_runs[max(longest_runs)][0]
            self.columns = [str(val) for val in self.loc[first_longest_index]]
            self.drop(self.loc[:first_longest_index].index.tolist(), axis=0, inplace=True)

    def _infer_range(self, mode: str = None) -> None:
        if mode is None:
            pass
        elif mode == Frame.InferRange.TRIM_SURROUNDING:
            self._trim_nulls_around_table()
        elif mode == Frame.InferRange.STRIP_NULLS:
            self._strip_fully_null()
        elif mode == Frame.InferRange.SMALLEST_VALID:
            self._truncate_after_valid()
        else:
            Frame.InferRange.raise_if_not_a_member(mode)

    def _trim_nulls_around_table(self) -> None:
        self._drop_rows_around_table()
        self._drop_columns_around_table()

    def _strip_fully_null(self) -> None:
        self._drop_fully_null_rows()
        self._drop_fully_null_columns()

    def _truncate_after_valid(self) -> None:
        self._drop_rows_around_table()
        self._truncate_rows_after_valid()

        self._drop_columns_around_table()
        self._truncate_columns_after_valid()

    def _drop_rows_around_table(self) -> None:
        nulls = [(index, row.isnull().all()) for index, row in self.iterrows()]
        invalid_rows: list = sum([list(itertools.takewhile(lambda row: row[1], iterable)) for iterable in (nulls, reversed(nulls))], [])
        if invalid_rows:
            self.drop([index for index, isnull in invalid_rows], axis=0, inplace=True)

    def _drop_columns_around_table(self) -> None:
        nulls = [(name, col.isnull().all()) for name, col in self.iteritems()]
        invalid_cols: list = sum([list(itertools.takewhile(lambda col: self._value_is_null(col[0]) and col[1], iterable)) for iterable in (nulls, reversed(nulls))], [])
        if invalid_cols:
            self.drop([name for name, isnull in invalid_cols], axis=1, inplace=True)

    def _truncate_rows_after_valid(self) -> None:
        empty_rows = self[self.isnull().all(axis=1) == True].index.tolist()
        if empty_rows:
            self.drop(self.loc[empty_rows[0]:].index.tolist(), axis=0, inplace=True)

    def _truncate_columns_after_valid(self) -> None:
        invalid_columns = list(itertools.dropwhile(lambda val: not self._value_is_null(val), self.columns))
        if invalid_columns:
            self.drop(invalid_columns, axis=1, inplace=True)

    def _drop_fully_null_rows(self) -> None:
        self.dropna(axis=0, how="all", inplace=True)

    def _drop_fully_null_columns(self) -> None:
        self.drop([name for name, col in self.iteritems() if self._value_is_null(name) and col.isnull().all()], axis=1, inplace=True)

    def _clean_dtypes(self) -> Frame:
        self._clean_nullable_ints()
        return self

    def _clean_nullable_ints(self) -> None:
        for name, col in self.iteritems():
            if col.dtype.name == "float64":
                if col.apply(lambda val: val is None or np.isnan(val) or val.is_integer()).all():
                    self[name] = col.astype("Int64")

    def _infer_boolean_columns(self) -> None:
        for name, col in self.iteritems():
            try:
                if col.isin([True, False, None]).all():
                    unfinished = col.replace(1, True).replace(0, False)
                    self[name] = unfinished.where(unfinished.notnull(), None)
            except (TypeError, SystemError):
                pass

    def _write_to_excel(self, writer: ExcelWriter, sheet_name: str, index: bool, **kwargs: Any) -> None:
        df, sheet_name = self.infer_dtypes(), Maybe(sheet_name).else_(self.DEFAULT_SHEET_NAME)
        super(type(df), df).to_excel(writer.writer, sheet_name=sheet_name, index=index, **kwargs)
        df._table_from_sheet(writer[sheet_name])

    def _table_from_sheet(self, sheet: Any) -> None:
        for idx, col in enumerate(self):
            series = self[col]
            datalen, headerlen = series.astype(str).map(len).max(), len(str(series.name))
            max_len = datalen + 1 if datalen > headerlen + 3 else headerlen + 4
            sheet.set_column(idx, idx, max_len)
        sheet.add_table(0, 0, len(sheet.table) - 1, len(sheet.table[0]) - 1, {"columns": [{"header": str(name)} for name in self.columns], "style": "Table Style Light 9"})

    @staticmethod
    def _value_is_null(value: Any) -> bool:
        return any([str(value).lower().startswith("unnamed"), *[str(value).lower() == nullproxy for nullproxy in ("none", "nan", "null", "")]])

    @staticmethod
    def _decrypted_stream_from_xlsx(path: PathLike, password: str) -> io.BytesIO:
        import msoffcrypto

        with open(path, "rb") as source_stream:
            file = msoffcrypto.OfficeFile(source_stream)
            file.load_key(password=password)
            destination_stream = io.BytesIO()
            file.decrypt(destination_stream)

        return destination_stream


class ExcelWriter:
    def __init__(self, filepath: PathLike) -> None:
        self.writer = pd.ExcelWriter(os.fspath(filepath), engine="xlsxwriter", mode="w")

    def __enter__(self) -> ExcelWriter:
        return self

    def __exit__(self, ex_type: Any, ex_value: Any, ex_traceback: Any) -> None:
        self.writer.save()

    def __getitem__(self, sheet_name: str) -> Any:
        # if self.writer.sheets is None:
        # self.writer.add_worksheet(sheet_name)
        return self.writer.sheets.get(sheet_name)
