from __future__ import annotations

from collections.abc import Mapping
from collections.abc import Sequence
from collections.abc import Set as AbstractSet
from pathlib import Path
from typing import Any
from typing import Literal
from typing import cast

from fastparquet import write
from hypothesis import assume
from hypothesis import given
from hypothesis.extra.pandas import column
from hypothesis.extra.pandas import data_frames
from hypothesis.extra.pandas import range_indexes
from hypothesis.strategies import DataObject
from hypothesis.strategies import SearchStrategy
from hypothesis.strategies import booleans
from hypothesis.strategies import data
from hypothesis.strategies import dictionaries
from hypothesis.strategies import floats
from hypothesis.strategies import integers
from hypothesis.strategies import lists
from hypothesis.strategies import none
from hypothesis.strategies import sampled_from
from hypothesis.strategies import sets
from hypothesis.strategies import tuples
from numpy import float32
from numpy import nan
from pandas import DataFrame
from pandas import RangeIndex
from pandas import Series
from pandas import concat
from pandas.testing import assert_frame_equal
from pandas.testing import assert_series_equal
from pytest import mark
from pytest import param
from pytest import raises

from utilities.fastparquet import _PARQUET_DTYPES
from utilities.fastparquet import EmptyDataFrameError
from utilities.fastparquet import InvalidDTypeError
from utilities.fastparquet import InvalidRowGroupIndexError
from utilities.fastparquet import _get_parquet_file
from utilities.fastparquet import count_rows
from utilities.fastparquet import get_columns
from utilities.fastparquet import get_dtypes
from utilities.fastparquet import get_num_row_groups
from utilities.fastparquet import read_parquet
from utilities.fastparquet import write_parquet
from utilities.hypothesis import dates_pd
from utilities.hypothesis import lists_fixed_length
from utilities.hypothesis import temp_paths
from utilities.hypothesis import text_ascii
from utilities.numpy import datetime64ns
from utilities.pandas import DataFrameRangeIndexError
from utilities.pandas import Int64
from utilities.pandas import string


class TestCountRows:
    @given(rows=lists(booleans(), min_size=1), root=temp_paths())
    def test_main(self, rows: Sequence[bool], root: Path) -> None:
        n = len(rows)
        df = DataFrame(rows, index=RangeIndex(n), columns=["value"])
        write_parquet(df, path := root.joinpath("df.parq"))
        result = count_rows(path)
        assert result == n


class TestGetColumns:
    @given(columns=sets(text_ascii()), root=temp_paths())
    def test_main(self, columns: AbstractSet[str], root: Path) -> None:
        as_list = list(columns)
        df = DataFrame(nan, index=RangeIndex(1), columns=as_list, dtype=float)
        write_parquet(df, path := root.joinpath("df.parq"))
        result = get_columns(path)
        assert result == as_list


class TestGetDtypes:
    @given(
        dtypes=dictionaries(
            text_ascii(), sampled_from(sorted(_PARQUET_DTYPES, key=repr))
        ),
        root=temp_paths(),
    )
    def test_main(self, dtypes: Mapping[str, Any], root: Path) -> None:
        df = DataFrame(None, index=RangeIndex(1), columns=list(dtypes)).astype(
            dtypes
        )
        write_parquet(df, path := root.joinpath("df.parq"))
        result = get_dtypes(path)
        assert result == dtypes


class TestGetNumRowGroups:
    def test_main(self, tmp_path: Path) -> None:
        df = DataFrame(0.0, index=RangeIndex(2), columns=["value"])
        write_parquet(
            df, path := tmp_path.joinpath("df.parq"), row_group_offsets=1
        )
        result = get_num_row_groups(path)
        expected = 2
        assert result == expected


class TestGetParquetFile:
    def test_main(self, tmp_path: Path) -> None:
        path = tmp_path.joinpath("file")
        df = DataFrame(0.0, index=RangeIndex(1), columns=["value"])
        write(path, df)
        _ = _get_parquet_file(path)

    def test_get_row_group(self, tmp_path: Path) -> None:
        path = tmp_path.joinpath("file")
        df = DataFrame(0.0, index=RangeIndex(2), columns=["value"])
        write(path, df, row_group_offsets=1)
        _ = _get_parquet_file(path, row_group=0)
        _ = _get_parquet_file(path, row_group=1)
        with raises(InvalidRowGroupIndexError):
            _ = _get_parquet_file(path, row_group=2)

    @mark.parametrize("as_str", [param(True), param(False)])
    def test_error(self, *, tmp_path: Path, as_str: bool) -> None:
        path = tmp_path.joinpath("file")
        path_use = path.as_posix() if as_str else path
        with raises(FileNotFoundError):
            _ = _get_parquet_file(path_use)


class TestReadAndWriteParquet:
    @given(data=data(), root=temp_paths(), as_series=booleans())
    @mark.parametrize(
        ("elements", "dtype"),
        [
            param(booleans(), bool),
            param(dates_pd() | none(), datetime64ns),
            param(floats(-10.0, 10.0) | none(), float),
            param(integers(-10, 10) | none(), Int64),
            param(text_ascii() | none(), string),
        ],
    )
    def test_writing_df(
        self,
        *,
        data: DataObject,
        elements: SearchStrategy[Any],
        dtype: Any,
        root: Path,
        as_series: bool,
    ) -> None:
        rows = data.draw(lists(elements, min_size=1))
        n = len(rows)
        df = DataFrame(rows, index=RangeIndex(n), columns=["value"]).astype(
            dtype
        )
        write_parquet(df, path := root.joinpath("df.parq"))
        head = data.draw(sampled_from([n, None]))
        columns = "value" if as_series else None
        read = read_parquet(path, head=head, columns=columns)
        if as_series:
            assert isinstance(read, Series)
            assert_series_equal(read, df["value"])
        else:
            assert isinstance(read, DataFrame)
            assert_frame_equal(read, df)

    @given(value1=floats(), value2=floats(), root=temp_paths())
    def test_writing_iterable_of_dfs(
        self, value1: float, value2: float, root: Path
    ) -> None:
        df = DataFrame(
            [[value1], [value2]], index=RangeIndex(2), columns=["value"]
        )
        parts = [df.iloc[:1], df.iloc[1:].reset_index(drop=True)]
        write_parquet(parts, path := root.joinpath("df.parq"))
        result = read_parquet(path)
        assert_frame_equal(result, df)

    @given(
        data=data(),
        column1=text_ascii(),
        column2=text_ascii(),
        root=temp_paths(),
    )
    def test_series_from_dataframe_with_two_string_columns(
        self, data: DataObject, column1: str, column2: str, root: Path
    ) -> None:
        _ = assume(column1 != column2)
        elements = text_ascii() | none()
        rows = data.draw(lists(tuples(elements, elements), min_size=1))
        df = DataFrame(
            rows,
            index=RangeIndex(len(rows)),
            columns=[column1, column2],
            dtype=string,
        )
        write_parquet(df, path := root.joinpath("df.parq"))
        sr = read_parquet(path, columns=column1)
        assert_series_equal(sr, df[column1])

    @given(value1=floats(), value2=floats(), root=temp_paths())
    def test_read_row_groups(
        self, value1: float, value2: float, root: Path
    ) -> None:
        df = DataFrame(
            [[value1], [value2]], index=RangeIndex(2), columns=["value"]
        )
        write_parquet(df, path := root.joinpath("df.parq"), row_group_offsets=1)
        result1 = read_parquet(path, row_group=0)
        expected1 = df.iloc[:1]
        assert_frame_equal(result1, expected1)
        result2 = read_parquet(path, row_group=1)
        expected2 = df.iloc[1:].reset_index(drop=True)
        assert_frame_equal(result2, expected2)
        with raises(InvalidRowGroupIndexError):
            _ = read_parquet(path, row_group=2)

    @given(
        data=data(),
        root=temp_paths(),
        num_dfs=integers(1, 10),
        row_group_offsets=integers(1, 10),
    )
    def test_iterable_of_dfs_with_strings(
        self, data: DataObject, root: Path, num_dfs: int, row_group_offsets: int
    ) -> None:
        def as_str(df: DataFrame, /) -> DataFrame:
            return df.astype(cast(Literal["string"], string))

        elements = data_frames(
            [cast(Any, column)("value", elements=text_ascii())],
            index=range_indexes(min_size=1),
        ).map(as_str)
        dfs = data.draw(lists_fixed_length(elements, num_dfs))
        write_parquet(
            dfs,
            path := root.joinpath("df.parq"),
            row_group_offsets=row_group_offsets,
        )
        result = read_parquet(path)
        expected = concat(dfs).reset_index(drop=True)
        assert_frame_equal(result, expected)


class TestWriteParquet:
    @given(value=text_ascii(), root=temp_paths())
    def test_strings_are_stored_as_objects(
        self, value: str, root: Path
    ) -> None:
        df = DataFrame(
            value, index=RangeIndex(1), columns=["value"], dtype=string
        )
        write_parquet(df, path := root.joinpath("df.parq"))
        file = _get_parquet_file(path)
        dtypes = file.dtypes
        expected = {"value": object}
        assert dtypes == expected

    def test_empty_dataframe_error(self, tmp_path: Path) -> None:
        df = DataFrame()
        with raises(EmptyDataFrameError):
            write_parquet(df, tmp_path.joinpath("df.parq"))

    def test_check_range_index(self, tmp_path: Path) -> None:
        df = DataFrame(nan, index=[0], columns=["value"])
        with raises(DataFrameRangeIndexError):
            write_parquet(df, tmp_path.joinpath("df.parq"))

    def test_check_invalid_dtype(self, tmp_path: Path) -> None:
        df = DataFrame(
            nan, index=RangeIndex(1), columns=["value"], dtype=float32
        )
        with raises(InvalidDTypeError):
            write_parquet(df, tmp_path.joinpath("df.parq"))

    def test_extra_dtype(self, tmp_path: Path) -> None:
        df = DataFrame(
            nan, index=RangeIndex(1), columns=["value"], dtype=float32
        )
        write_parquet(
            df, tmp_path.joinpath("df.parq"), extra_dtypes={"value": float32}
        )

    def test_extra_dtype_ignored(self, tmp_path: Path) -> None:
        df = DataFrame(
            nan, index=RangeIndex(1), columns=["value"], dtype=float32
        )
        with raises(InvalidDTypeError):
            write_parquet(
                df,
                tmp_path.joinpath("df.parq"),
                extra_dtypes={"other": float32},
            )
