"""
Artifact Base module.

Defines the base classes for system artifacts.
"""

from __future__ import annotations

import asyncio
from typing import Any, Dict, Mapping, Optional, Type
from datetime import date, datetime

import pandas as pd
import dask.dataframe as dd

# Correct internal import for the base class
from .parquet import ParquetArtifact

DateLike = str | date | datetime | None


def _validate_and_format_date(name: str, value: DateLike) -> Optional[str]:
    """
    Normalize date-like input into a canonical string '%Y-%m-%d'.

    - None -> None
    - str/date/datetime -> parse with pandas.to_datetime, take .date(), return '%Y-%m-%d'
    - else -> raise TypeError
    """
    if value is None:
        return None
    if isinstance(value, (str, date, datetime)):
        try:
            return pd.to_datetime(value).date().strftime("%Y-%m-%d")
        except Exception as e:
            raise ValueError(f"{name} must be a valid date; got {value!r}") from e
    raise TypeError(f"{name} must be str, date, datetime, or None; got {type(value)}")


class BaseArtifact(ParquetArtifact):
    """
    Generic base class for Parquet-based artifacts with enhanced date handling and async support.

    Features:
    - Automatic date normalization to '%Y-%m-%d' strings.
    - Async loading support via `aload`.
    - Lifecycle hooks (before/after load).
    - Params serialization.

    The `config` class attribute can be used to set default configuration values.
    """

    config: Mapping[str, Any] = {}

    # Typed attributes for IDE support (populated in __init__)
    parquet_start_date: Optional[str]
    parquet_end_date: Optional[str]
    provider_class: Optional[Type[Any]]
    class_params: Dict[str, Any]
    df: Optional[pd.DataFrame | dd.DataFrame] = None

    def __init__(
        self,
        **kwargs: Any,
    ) -> None:
        # 1. Merge defaults from class config with runtime kwargs
        merged = {**self.config, **kwargs}

        # 2. Extract and normalize date fields *before* passing to super
        #    This ensures ParquetArtifact receives clean strings if it validates them.
        p_start = merged.get("parquet_start_date")
        p_end = merged.get("parquet_end_date")

        self.parquet_start_date = _validate_and_format_date(
            "parquet_start_date", p_start
        )
        self.parquet_end_date = _validate_and_format_date("parquet_end_date", p_end)

        # Update merged dict with normalized values to ensure consistency
        if self.parquet_start_date:
            merged["parquet_start_date"] = self.parquet_start_date
        if self.parquet_end_date:
            merged["parquet_end_date"] = self.parquet_end_date

        # 3. Initialize parent (ParquetArtifact handles config validation and resource setup)
        super().__init__(**merged)

        # 4. Store additional wrapper config
        self.provider_class = merged.get("provider_class", None)

        # Default class_params from instance attributes if not provided
        self.class_params = merged.get("class_params", None) or {
            "debug": self.debug,
            "logger": self.logger,
            "fs": self.fs,
            "verbose": getattr(self, "verbose", False),
        }

        # 5. Validation Logic
        if self.parquet_start_date and self.parquet_end_date:
            if self.parquet_start_date > self.parquet_end_date:
                raise ValueError(
                    f"parquet_start_date {self.parquet_start_date} "
                    f"cannot be after parquet_end_date {self.parquet_end_date}"
                )

    # -------- Lifecycle Hooks (No-op defaults) --------

    def before_load(self, **kwargs: Any) -> None:
        """Hook called synchronously before loading data."""
        return None

    def after_load(self, **kwargs: Any) -> None:
        """Hook called synchronously after loading data."""
        return None

    async def abefore_load(self, **kwargs: Any) -> None:
        """Hook called asynchronously before loading data."""
        return None

    async def aafter_load(self, **kwargs: Any) -> None:
        """Hook called asynchronously after loading data."""
        return None

    # -------- Public API --------

    def load(self, **kwargs: Any) -> pd.DataFrame | dd.DataFrame:
        """
        Synchronous load with hooks.
        Delegates actual loading to super().load().
        """
        self.before_load(**kwargs)
        self.df = super().load(**kwargs)
        self.after_load(**kwargs)
        return self.df

    async def aload(self, **kwargs: Any) -> pd.DataFrame | dd.DataFrame:
        """
        Asynchronous load with hooks.
        Runs the blocking super().load() in a thread.
        """
        await self.abefore_load(**kwargs)
        # Run blocking IO in a separate thread to avoid blocking the event loop
        df = await asyncio.to_thread(super().load, **kwargs)
        self.df = df
        await self.aafter_load(**kwargs)
        return self.df

    def has_date_window(self) -> bool:
        """Check if a date window is defined."""
        return bool(self.parquet_start_date or self.parquet_end_date)

    def date_window(self) -> tuple[Optional[str], Optional[str]]:
        """Return the (start, end) date window tuple."""
        return self.parquet_start_date, self.parquet_end_date

    def to_params(self) -> Dict[str, Any]:
        """Serialize configuration parameters."""
        return {
            "parquet_start_date": self.parquet_start_date,
            "parquet_end_date": self.parquet_end_date,
            "provider_class": self.provider_class,
            "class_params": dict(self.class_params),
            # Add useful metadata from parent
            "parquet_storage_path": getattr(self, "_storage_path", None),
        }
