import datetime as dt
import posixpath
from typing import Optional, List, Any

import dask.dataframe as dd
import pandas as pd
import pyarrow.dataset as ds
import pyarrow as pa
from pydantic import BaseModel, ConfigDict, model_validator

from sibi_flux.logger import Logger


class ParquetConfig(BaseModel):
    """
    Universal configuration for loading Parquet data.
    Supports:
      1. Single File (parquet_filename set)
      2. Partitioned Range (parquet_start/end_date set)
      3. Generic Folder Scan (No dates, no filename -> loads everything)
    """

    # ---- Inputs ----
    parquet_storage_path: str
    parquet_filename: Optional[str] = None
    parquet_start_date: Optional[str] = None
    parquet_end_date: Optional[str] = None
    parquet_max_age_minutes: int = 0

    fs: Any
    logger: Optional[Logger] = None
    debug: bool = False

    # Optional: If set, forces Hive logic. If None, auto-detects.
    partition_on: Optional[list[str]] = None

    # ---- Derived ----
    parquet_full_path: Optional[str] = None
    parquet_is_recent: bool = False
    load_parquet: bool = False

    model_config = ConfigDict(arbitrary_types_allowed=True)

    @model_validator(mode="after")
    def _normalize(self):
        if self.logger is None:
            self.logger = Logger.default_logger(logger_name=self.__class__.__name__)

        # 1. Path Check
        if not self.parquet_storage_path:
            raise ValueError("Parquet storage path must be specified")
        self.parquet_storage_path = self.parquet_storage_path.rstrip("/")

        # 2. Date Presence Check
        if bool(self.parquet_start_date) != bool(self.parquet_end_date):
            raise ValueError(
                "Both start_date and end_date must be provided, or neither."
            )

        # 3. Date Order Check (FIX: Add this block)
        if self.parquet_start_date and self.parquet_end_date:
            try:
                start = dt.datetime.strptime(self.parquet_start_date, "%Y-%m-%d").date()
                end = dt.datetime.strptime(self.parquet_end_date, "%Y-%m-%d").date()
                if end < start:
                    index_cols: list[str | None] = []
                    raise ValueError("end_date cannot be before start_date")
            except ValueError as e:
                # Re-raise explicit logic errors, ignore parsing errors (let Pydantic handle type checks)
                if "end_date cannot be before" in str(e):
                    raise e

        # 4. Path Resolution
        if self.parquet_filename:
            raw = posixpath.join(self.parquet_storage_path, self.parquet_filename)
            self.parquet_full_path = self.ensure_file_extension(raw, "parquet")
        else:
            self.parquet_full_path = self.parquet_storage_path

        return self

    def load_files(self, filters: Optional[List[Any]] = None) -> dd.DataFrame:
        """
        Smart loader that decides the best strategy based on config.
        """
        # Strategy 1: Specific File (Fastest)
        if self.parquet_filename:
            if not self.determine_recency():
                return self._empty_ddf()
            # Explicitly annotate to handle list invariance with Optional[str]
            # Explicitly annotate to handle list invariance with Optional[str]
            files_to_load: List[str] = (
                [self.parquet_full_path] if self.parquet_full_path else []
            )

        # Strategy 2: Date Partitioned (Optimized Scan)
        elif self.parquet_start_date and self.parquet_end_date:
            files_to_load = self._discover_partitioned_files()

        # Strategy 3: Generic Folder Scan (Fallback)
        else:
            if self.logger:
                self.logger.debug(
                    f"No dates or filename provided. Scanning all files in {self.parquet_full_path}"
                )
            files_to_load = self._discover_all_files()

        if not files_to_load:
            if self.logger:
                self.logger.warning("No files found.")
            return self._empty_ddf()

        if self.logger:
            self.logger.debug(
                f"Loading {files_to_load} from {self.parquet_storage_path}"
            )

        self.load_parquet = True
        try:
            #
            # Found the exact files needed (files_to_load). Metadata file may be stale.
            # Don't waste time looking for them (filesystem, ignore_metadata).
            # Use the fast C++ parser (pyarrow).
            # If the files are tiny, bundle them (aggregate_files),
            # Use the file footers to skip data I don't need (filters).
            #
            return dd.read_parquet(
                files_to_load,
                engine="pyarrow",
                filesystem=self.fs,
                filters=filters,
                aggregate_files=True,
                ignore_metadata_file=True,  # we use "files_to_load" directly.  Metadata file may be stale.
            )
        except Exception as e:
            if self.logger:
                self.logger.error("Parquet error: %s", e)
            raise e

    # ------------------------- Discovery Logic -------------------------

    def _discover_partitioned_files(self) -> List[str]:
        # 1. Clean Path (Strip s3://)
        source_path = self.parquet_storage_path
        if "://" in source_path:
            source_path = source_path.split("://", 1)[1]

        # 2. Setup Partitioning
        if self.partition_on:
            partitioning = ds.partitioning(flavor="hive")
            key = self.partition_on[0]
        else:
            partitioning = ds.partitioning(
                schema=pa.schema(
                    [("year", pa.int32()), ("month", pa.int32()), ("day", pa.int32())]
                ),
                flavor=None,
            )
            key = "year"

        # 3. Create Dataset (Lazy Scan)
        try:
            dataset = ds.dataset(
                source=source_path,
                filesystem=self.fs,
                format="parquet",
                partitioning=partitioning,
            )
        except FileNotFoundError:
            # If the ROOT path itself doesn't exist, we return empty list safely
            return []

        # 4. Filter (Push-down)
        sd = dt.datetime.strptime(self.parquet_start_date, "%Y-%m-%d").date()
        ed = dt.datetime.strptime(self.parquet_end_date, "%Y-%m-%d").date()

        if self.partition_on:
            # String comparison for Hive
            expr = (ds.field(key) >= str(sd)) & (ds.field(key) <= str(ed))
        else:
            # Int comparison for Directory
            expr = (ds.field("year") >= sd.year) & (ds.field("year") <= ed.year)

        # 5. Materialize Files
        # get_fragments() ONLY returns files that exist on disk.
        # Non-existent or empty partition folders yield 0 fragments.
        fragments = dataset.get_fragments(expr)
        found_files = [f.path for f in fragments]

        if self.debug:
            if self.logger:
                self.logger.debug(f"Requested range: {sd} to {ed}")
                self.logger.debug(f"Actual files found: {len(found_files)}")
            # If 0 files found, it means the dates requested don't exist in storage

        return found_files

    def _discover_all_files(self) -> List[str]:
        """
        Scans EVERYTHING in the folder recursively.
        Used when no dates/partitions are specified.
        """
        try:
            # PyArrow Dataset without partitioning = Recursive Scan
            dataset = ds.dataset(
                self.parquet_storage_path,
                filesystem=self.fs,
                format="parquet",
                partitioning=None,  # Important: Treat everything as data, ignore folder structure
            )
            return dataset.files
        except Exception as e:
            if self.logger:
                self.logger.warning("Failed to scan folder: %s", e)
            return []

    # ------------------------- Helpers -------------------------

    def determine_recency(self) -> bool:
        path = self.parquet_full_path
        if not path or not path.endswith(".parquet"):
            return True

        if hasattr(self.fs, "get_file_info"):
            info = self.fs.get_file_info(path)
            if info.type == pa.fs.FileType.NotFound:
                return False
            mdt = info.mtime
        else:
            if not self._exists(path):
                return False
            mdt = self._get_mtime_fallback(path)

        if self.parquet_max_age_minutes == 0:
            return True
        if not mdt:
            return False

        now = dt.datetime.now(dt.timezone.utc)
        if mdt.tzinfo is None:
            mdt = mdt.replace(tzinfo=dt.timezone.utc)

        return (now - mdt) <= dt.timedelta(minutes=self.parquet_max_age_minutes)

    def _get_mtime_fallback(self, path):
        # (Legacy fsspec support)
        try:
            info = self.fs.info(path)
            # ... (timestamp conversion logic)
            return info["mtime"]
        except Exception:
            return None

    def _exists(self, path):
        try:
            return self.fs.exists(path)
        except Exception:
            return False

    @staticmethod
    def _empty_ddf():
        return dd.from_pandas(pd.DataFrame(), npartitions=1)

    @staticmethod
    def ensure_file_extension(filepath: str, extension: str) -> str:
        if filepath.endswith(f".{extension}"):
            return filepath
        return f"{filepath}.{extension}"
