"""
Parquet Artifact module.

Defines the ParquetArtifact class for orchestrating Parquet dataset generation.
"""

from __future__ import annotations

import datetime as dt
import threading
from functools import cached_property
from typing import Any, Dict, Type, TypeVar, Optional, Set, List

from pydantic import BaseModel, Field, ConfigDict, ValidationError

from sibi_flux.core import SecureResource
from sibi_flux.utils.date_utils import DateUtils

from sibi_flux.artifacts.parquet_engine import (
    Executor as DataWrapper,
    UpdatePlanner,
    MissingManifestManager,
)

T = TypeVar("T")


class ArtifactConfig(BaseModel):
    """
    Strict configuration schema for ParquetArtifact.
    Validates arguments immediately upon initialization.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")

    parquet_storage_path: str
    parquet_filename: str

    provider_class: Optional[Type] = None
    partition_on: List[str] = Field(default_factory=lambda: ["partition_date"])
    hive_style: bool = True
    overwrite: bool = False

    show_progress: bool = False
    timeout: float = 30.0
    max_threads: int = 3

    max_retries: int = 3
    backoff_base: float = 2.0

    load_params: Dict[str, Any] = Field(default_factory=dict)


class ParquetArtifact(SecureResource):
    """
    Orchestrates the generation of a single date-partitioned Parquet dataset.

    Features:
      - Validated configuration via Pydantic.
      - Thread-safe concurrency control (process-local).
      - Lazy initialization of resources (manifests, planners).
      - Safe resource cleanup.
    """

    _global_lock = threading.RLock()
    _active_runs: Set[tuple[str, str]] = set()
    logger_extra = {"sibi_flux_component": __name__}

    def __init__(self, **kwargs: Any):
        """
        Initializes the ParquetArtifact.

        Args:
            **kwargs: Configuration parameters validated against ArtifactConfig.
        """
        try:
            self.config = ArtifactConfig(**kwargs)
        except ValidationError as e:
            raise ValueError(f"Invalid ParquetArtifact configuration: {e}") from e

        self.all_kwargs = self.config.model_dump()

        super().__init__(**self.all_kwargs)

        self._lock = threading.RLock()

        self._storage_path = self.config.parquet_storage_path
        self._provider_class = self.config.provider_class

        self.logger_extra.update(
            {
                "artifact_storage_path": self._storage_path,
            }
        )

    def _invalidate_cached(self, *names: str) -> None:
        """
        Invalidate cached properties by name.

        SAFEGUARD: Tries to call .save() on objects (like manifests)
        before removing them from memory to prevent state loss.
        """
        for name in names:
            if name in self.__dict__:
                obj = self.__dict__[name]

                if hasattr(obj, "save") and callable(obj.save):
                    try:
                        self.logger.debug(
                            f"Saving state for '{name}' before cache invalidation...",
                            extra=self.logger_extra,
                        )
                        obj.save()
                    except Exception as e:
                        self.logger.warning(
                            f"Failed to save '{name}' state during invalidation: {e}",
                            extra=self.logger_extra,
                        )

                del self.__dict__[name]

    def _build_manifest_path(self) -> str:
        """Constructs the path for the missing data manifest."""
        base = self._storage_path.rstrip("/") + "/"
        return f"{base}_manifests/missing.parquet"

    @cached_property
    def mmanifest(self) -> MissingManifestManager:
        self.logger.debug(
            "Initializing MissingManifestManager...", extra=self.logger_extra
        )
        manifest_path = self._build_manifest_path()

        manifest_dir = (
            manifest_path.rsplit("/", 1)[0] if "/" in manifest_path else manifest_path
        )
        self.ensure_directory_exists(manifest_dir)

        mgr = MissingManifestManager(
            fs=self.fs,
            manifest_path=manifest_path,
            clear_existing=self.config.overwrite,
            debug=self.debug,
            logger=self.logger,
            overwrite=self.config.overwrite,
        )

        if hasattr(mgr, "_safe_exists") and not mgr._safe_exists(mgr.manifest_path):
            self.logger.debug(
                f"Creating new manifest at {mgr.manifest_path}", extra=self.logger_extra
            )
            mgr.save()

        return mgr

    @cached_property
    def update_planner(self) -> UpdatePlanner:
        self.logger.debug("Initializing UpdatePlanner...", extra=self.logger_extra)
        skipped_files: List[Any] = self.mmanifest.load_existing() or []

        cfg = {
            **self.all_kwargs,
            "fs": self.fs,
            "debug": self.debug,
            "logger": self.logger,
            "description": getattr(self._provider_class, "__name__", "DataWrapper"),
            "skipped": list(skipped_files),
            "mmanifest": self.mmanifest,
            "partition_on": self.config.partition_on,
            "hive_style": self.config.hive_style,
        }
        return UpdatePlanner(**cfg)

    @cached_property
    def data_wrapper(self) -> DataWrapper:
        self.logger.debug("Initializing DataWrapper...", extra=self.logger_extra)

        _ = self.update_planner.plan

        class_params = {
            "debug": self.debug,
            "logger": self.logger,
            "fs": self.fs,
            "verbose": self.verbose,
        }

        cfg = {
            "data_path": self._storage_path,
            "fs": self.fs,
            "debug": self.debug,
            "logger": self.logger,
            "verbose": self.verbose,
            "provider_class": self._provider_class,
            "class_params": class_params,
            "load_params": self.config.load_params,
            "mmanifest": self.mmanifest,
            "update_planner": self.update_planner,
            "date_field": self.all_kwargs.get("date_field"),
            "show_progress": self.config.show_progress,
            "timeout": self.config.timeout,
            "max_threads": self.config.max_threads,
        }
        return DataWrapper(**cfg)

    def load(self, **kwargs: Any):
        """
        Directly load data using the configured provider_class.
        Bypasses planning/manifest process.
        """
        if not self._provider_class:
            raise ValueError("provider_class is not configured.")

        params = {
            "backend": "parquet",
            "fs": self.fs,
            "logger": self.logger,
            "debug": self.debug,
            "parquet_storage_path": self._storage_path,
            "parquet_start_date": self.all_kwargs.get("parquet_start_date"),
            "parquet_end_date": self.all_kwargs.get("parquet_end_date"),
            "partition_on": self.config.partition_on,
            **(self.all_kwargs.get("class_params") or {}),
        }

        with self._provider_class(**params) as instance:
            return instance.load(**kwargs)

    def generate_parquet(self, **kwargs: Any) -> None:
        """
        Generate or update the Parquet dataset according to the plan.
        - Merges runtime kwargs.
        - Invalidates dependent cached properties safely.
        - Guards against duplicate concurrent runs.
        """
        self.all_kwargs.update(kwargs)

        self._invalidate_cached("update_planner", "data_wrapper")
        if "overwrite" in kwargs and kwargs["overwrite"]:
            self._invalidate_cached("mmanifest")

        key = self._storage_path
        with ParquetArtifact._global_lock:
            if key in ParquetArtifact._active_runs:
                self.logger.info(
                    f"Run active for {key}; skipping.", extra=self.logger_extra
                )
                return
            ParquetArtifact._active_runs.add(key)

        try:
            self.ensure_directory_exists(self._storage_path)

            plan = self.update_planner.plan
            if plan is None or (hasattr(plan, "empty") and plan.empty):
                self.logger.info("No updates needed.", extra=self.logger_extra)
                return

            if self.config.show_progress and not getattr(
                self.update_planner, "_printed_this_run", False
            ):
                try:
                    self.update_planner.show_update_plan()
                    setattr(self.update_planner, "_printed_this_run", True)
                except Exception:
                    pass

            dw_retry_kwargs = {
                "max_retries": self.config.max_retries,
                "backoff_base": self.config.backoff_base,
            }
            if "backoff_jitter" in self.all_kwargs:
                dw_retry_kwargs["backoff_jitter"] = self.all_kwargs["backoff_jitter"]
            if "backoff_max" in self.all_kwargs:
                dw_retry_kwargs["backoff_max"] = self.all_kwargs["backoff_max"]

            with self._lock:
                dw = self.data_wrapper
                if hasattr(dw, "process"):
                    dw.process(**dw_retry_kwargs)

                    if self.config.show_progress and hasattr(
                        dw, "show_benchmark_summary"
                    ):
                        dw.show_benchmark_summary()

        finally:
            with ParquetArtifact._global_lock:
                ParquetArtifact._active_runs.discard(key)

    def update_parquet(self, period: str = "today", **kwargs: Any) -> None:
        """
        High-level entry point to update Parquet for a standard or custom period.
        """
        final_kwargs = {**self.all_kwargs, **kwargs}
        period_params: Dict[str, dt.date] = {}

        if period == "itd":
            start_str = final_kwargs.get("history_begins_on")
            if not start_str:
                raise ValueError("Period 'itd' requires 'history_begins_on'.")
            period_params = {
                "parquet_start_date": dt.datetime.strptime(
                    start_str, "%Y-%m-%d"
                ).date(),
                "parquet_end_date": dt.date.today(),
            }
        elif period == "ytd":
            period_params = {
                "parquet_start_date": dt.date(dt.date.today().year, 1, 1),
                "parquet_end_date": dt.date.today(),
            }
        elif period == "custom":
            p = final_kwargs
            s = p.get("start_on") or p.get("start_date") or p.get("start")
            e = p.get("end_on") or p.get("end_date") or p.get("end")

            if not s or not e:
                raise ValueError("Period 'custom' requires 'start_on' and 'end_on'.")

            period_params = {
                "parquet_start_date": dt.datetime.strptime(str(s), "%Y-%m-%d").date(),
                "parquet_end_date": dt.datetime.strptime(str(e), "%Y-%m-%d").date(),
            }
        else:
            try:
                s, e = DateUtils.parse_period(period=period)
                period_params = {"parquet_start_date": s, "parquet_end_date": e}
            except Exception as e:
                raise ValueError(f"Failed to parse period '{period}': {e}") from e

        final_kwargs.update(period_params)
        self.generate_parquet(**final_kwargs)

    def ensure_directory_exists(self, path: str) -> None:
        """Robust directory creation handling various fs backends."""
        with self._lock:
            if not self.fs.exists(path):
                try:
                    self.fs.makedirs(path, exist_ok=True)
                except TypeError:
                    # Fallback for backends not supporting exist_ok
                    try:
                        self.fs.makedirs(path)
                    except FileExistsError:
                        pass

    def _cleanup(self) -> None:
        """
        Cleanup resources.
        Calls _invalidate_cached to ensure manifests are saved before closing.
        """
        try:
            self._invalidate_cached("mmanifest")

            if "data_wrapper" in self.__dict__ and hasattr(self.data_wrapper, "close"):
                self.data_wrapper.close()
        except Exception as e:
            self.logger.warning(
                f"Error during ParquetArtifact cleanup: {e}", extra=self.logger_extra
            )
