"""
Parquet Readers module.

Provides specialized readers for Parquet files handling partitioning and types.
"""

from __future__ import annotations

from collections.abc import Mapping
from typing import (
    TYPE_CHECKING,
    Any,
    ClassVar,
)

import dask.dataframe as dd
import pandas as pd

from sibi_flux.df_helper._df_helper import DfHelper

if TYPE_CHECKING:
    import fsspec


class ParquetReader(DfHelper):
    """
    This class is a specialised helper for reading and managing Parquet files.

    The `ParquetReader` class is designed to facilitate working with Parquet
    datasets stored across different filesystems. It initialises the required
    resources, ensures the existence of the specified Parquet directory,
    and provides an abstraction to load the data into a Dask DataFrame.
    """

    DEFAULT_CONFIG: ClassVar[Mapping[str, Any]] = {
        "backend": "parquet",
        "partition_on": ["partition_date"],
    }
    df: pd.DataFrame | dd.DataFrame | None = None

    def __init__(
        self,
        parquet_start_date: str | None = None,
        parquet_end_date: str | None = None,
        **kwargs: Any,
    ) -> None:
        # Merge explicitly provided dates into kwargs config
        if parquet_start_date:
            kwargs["parquet_start_date"] = parquet_start_date
        if parquet_end_date:
            kwargs["parquet_end_date"] = parquet_end_date

        self.config = {
            **self.DEFAULT_CONFIG,
            **kwargs,
        }
        super().__init__(**self.config)

        # Access validated properties from self.backend_params (the Pydantic model)
        # instead of the raw config dict, ensuring type safety.
        if hasattr(self, "backend_params"):
            self.parquet_storage_path: str = self.backend_params.parquet_storage_path
            self.parquet_start_date: str = self.backend_params.parquet_start_date
            self.parquet_end_date: str = self.backend_params.parquet_end_date
            self.fs: fsspec.AbstractFileSystem = self.backend_params.fs
        else:
            # Fallback if DfHelper structure differs
            self.parquet_storage_path = self.config.get("parquet_storage_path")
            self.parquet_start_date = self.config.get("parquet_start_date")
            self.parquet_end_date = self.config.get("parquet_end_date")
            self.fs = self.config.get("fs")

        # Additional filesystem checks can remain
        if not self.directory_exists():
            raise ValueError(f"{self.parquet_storage_path} does not exist")

    def load(self, **kwargs: Any) -> pd.DataFrame | dd.DataFrame:
        self.df = super().load(**kwargs)
        return self.df

    async def aload(self, **kwargs: Any) -> pd.DataFrame | dd.DataFrame:
        self.df = await super().aload(**kwargs)
        return self.df

    def directory_exists(self) -> bool:
        try:
            info = self.fs.info(self.parquet_storage_path)
            return info["type"] == "directory"
        except (FileNotFoundError, AttributeError):
            return False
