from __future__ import annotations

import datetime as dt
import re
from typing import Iterable, Any

import fsspec
import pyarrow.dataset as ds
from fsspec.utils import infer_storage_options
from sibi_flux.logger import Logger


class FilePathGenerator:
    """
    Scans Hive-partitioned directories (key=value) and returns paths for a specific date range.
    """

    def __init__(
        self,
        base_path: str = "",
        *,
        fs: fsspec.AbstractFileSystem | None = None,
        date_partition_key: str = "partition_date",
        logger: Logger | None = None,
        debug: bool = False,
        storage_options: dict[str, Any] | None = None,
        exclude_patterns: Iterable[str] | None = None,
        file_extension: str = "parquet",
    ) -> None:
        self.logger = logger or Logger.default_logger(
            logger_name=self.__class__.__name__
        )
        self.debug = debug
        self.storage_options = storage_options or {}
        self.file_extension = file_extension.lstrip(".")
        self._compiled_exclusions = [re.compile(p) for p in (exclude_patterns or [])]
        self.date_partition_key = date_partition_key

        # 1. Setup Filesystem (Protocol Stripping)
        if fs:
            self.fs = fs
            self.base_path = infer_storage_options(base_path)["path"]
        else:
            self.fs, self.base_path = fsspec.core.url_to_fs(
                base_path, **self.storage_options
            )

        # 2. Strict Hive Partitioning
        # This automatically parses "key=value" directory names.
        self.partitioning = ds.partitioning(flavor="hive")

    def generate_file_paths(
        self, start_date: str | dt.date, end_date: str | dt.date, engine: str = "dask"
    ) -> list[str]:
        sd = self._to_date(start_date)
        ed = self._to_date(end_date)
        if sd > ed:
            sd, ed = ed, sd

        # Create Dataset
        try:
            dataset = ds.dataset(
                source=self.base_path,
                filesystem=self.fs,
                format=self.file_extension,
                partitioning=self.partitioning,
            )
        except Exception as e:
            if self.debug:
                self.logger.error(f"Failed to load dataset: {e}")
            return []

        # 3. Dynamic Filter Logic
        schema_names = set(dataset.schema.names)
        filter_expr = None
        filter_mode = "unknown"

        # Case A: User-specified single column (e.g., partition_date=2024-01-01)
        if self.date_partition_key in schema_names:
            filter_mode = "single_key"
            # We filter treating the key as a string comparable to YYYY-MM-DD
            filter_expr = (ds.field(self.date_partition_key) >= str(sd)) & (
                ds.field(self.date_partition_key) <= str(ed)
            )

        # Case B: Standard Composite Hierarchy (year=..., month=..., day=...)
        # We auto-fallback to this if the single key isn't found
        elif {"year", "month", "day"}.issubset(schema_names):
            filter_mode = "composite_ymd"
            # Broad filter on year first for performance
            filter_expr = (ds.field("year") >= sd.year) & (ds.field("year") <= ed.year)

        else:
            self.logger.warning(
                f"Could not find partition key in schema: {schema_names}."
            )

        # 4. Materialize & Refine
        # We assume PyArrow's filter is "coarse" (partition level), so we double-check checks in Python
        valid_paths = []
        fragments = dataset.get_fragments(filter=filter_expr)

        for frag in fragments:
            if self._is_file_in_range(frag, sd, ed, filter_mode):
                path = frag.path
                # Ensure full protocol paths for Pandas/Dask
                if not path.startswith(self.fs.protocol):
                    path = self.fs.unstrip_protocol(path)

                if not self._is_excluded(path):
                    valid_paths.append(path)

        return sorted(valid_paths)

    def _is_file_in_range(self, frag: Any, sd: dt.date, ed: dt.date, mode: str) -> bool:
        """ Parses partition keys from fragment metadata and validates exact date range. """
        try:
            keys = ds._get_partition_keys(frag.partition_expression)

            if mode == "single_key":
                val = keys.get(self.date_partition_key)
                if not val:
                    return False

                # Handle String vs Date type inference
                f_date = (
                    val
                    if isinstance(val, (dt.date, dt.datetime))
                    else self._to_date(val)
                )
                # Ensure comparison between dates (not datetime vs date)
                if isinstance(f_date, dt.datetime):
                    f_date = f_date.date()
                return sd <= f_date <= ed

            elif mode == "composite_ymd":
                f_date = dt.date(
                    int(keys["year"]), int(keys["month"]), int(keys["day"])
                )
                return sd <= f_date <= ed

            return True

        except Exception:
            # If parsing fails (e.g. malformed folder name), exclude safely
            return False

    def _is_excluded(self, path: str) -> bool:
        return any(pat.search(path) for pat in self._compiled_exclusions)

    @staticmethod
    def _to_date(x: Any) -> dt.date:
        if isinstance(x, dt.datetime):
            return x.date()
        if isinstance(x, dt.date):
            return x
        try:
            return dt.datetime.strptime(str(x), "%Y-%m-%d").date()
        except ValueError:
            # Fallback for weird formats if necessary
            return dt.date.min
