from __future__ import annotations

import datetime as dt
from collections.abc import Mapping, MutableMapping
from typing import Any, Optional

# Note: Ideally status tracking should be injected or imported from a shared location.
# For now, we assume LAST_STATUS is available or passed in.
# Since the user code had `from suite.etl.routers.status import LAST_STATUS`,
# we'll keep it or make it generic. To keep it generic in async_utils,
# we should probably accept a callback or status dict, but for now I'll retain the import
# if it exists, or mock it if this is a library file.
# However, `suite` implies application code.
# I will make the status update generic (optional callback) to avoid app-dependency in the lib.

from sibi_flux.dask_cluster.client_manager import DaskClientMixin
from sibi_flux.orchestration import ArtifactOrchestrator
from sibi_flux.logger import Logger
from sibi_flux.core import SecureResource


class PipelineExecutor(DaskClientMixin, SecureResource):
    """
    Orchestrator entry point for ETL updates.
    Manages Dask client lifecycle and delegates work to ArtifactOrchestrator.
    """

    def __init__(
        self,
        *,
        wrapped_classes: Mapping[str, list],
        logger: Logger,
        fs,
        debug: bool = False,
        verbose: bool = False,
        use_async_orchestrator: bool = True,
        artifact_class_kwargs: Optional[Mapping[str, Any]] = None,
        status_tracker: Optional[
            MutableMapping
        ] = None,  # Dependency injection for status
        **kwargs,
    ) -> None:
        ManagedResource.__init__(
            self, logger=logger, fs=fs, debug=debug, verbose=verbose
        )
        DaskClientMixin.__init__(self, logger=logger)

        # Dask client configuration
        self.max_workers = kwargs.get("max_workers", 8)
        self.max_threads = kwargs.get("max_threads", 10)
        dask_client = kwargs.get("dask_client")

        self._init_dask_client(
            dask_client=dask_client,
            logger=self.logger,
            n_workers=self.max_workers,
            threads_per_worker=self.max_threads,
            lazy=True,  # Do not block MainThread on init
        )

        # Core attributes
        self.wrapped_classes = wrapped_classes
        self._artifact_class_kwargs = artifact_class_kwargs or {}
        self._use_async_orchestrator = use_async_orchestrator
        self.status_tracker = status_tracker if status_tracker is not None else {}

        self.extra_logger = {
            "action_module_name": self.__class__.__name__,
            "start_date": dt.datetime.today().strftime("%Y-%m-%d"),
            "start_time": dt.datetime.today().strftime("%H:%M:%S"),
        }

        # Common orchestrator arguments
        self._orch_kwargs: MutableMapping[str, Any] = dict(
            wrapped_classes=self.wrapped_classes,
            max_workers=self.max_workers,
            artifact_class_kwargs=self._artifact_class_kwargs,
            logger=self.logger,
            fs=fs,
            debug=debug,
            verbose=verbose,
            dask_client=dask_client,
            use_dask=dask_client is not None,
        )

        self._orchestrator: Optional[ArtifactOrchestrator] = None

    # ---------------- Core Execution ----------------

    async def run(self, period: str, **params: Any) -> None:
        """Run ETL update for the given period."""
        self._validate_range_params_if_needed(period, params)

        log_extra = {**self.extra_logger, "period": period, "params": params}
        self.logger.info(f"Running update for period: {period}", extra=log_extra)

        # Initialize orchestrator only once per session
        if self._orchestrator is None:
            self._orchestrator = ArtifactOrchestrator(**self._orch_kwargs)

        results = await self._orchestrator.update_data(period, **params)
        self.logger.debug(
            f"Received {len(results)} computed results from orchestrator."
        )

        # Update status tracker if provided
        if self.status_tracker is not None:
            self.status_tracker.update(self._orchestrator.get_update_status())

    # ---------------- Convenience Shortcuts ----------------

    async def run_the_update_today(self, **params: Any) -> None:
        await self.run("today", **params)

    async def run_the_update_current_month(self, **params: Any) -> None:
        await self.run("current_month", **params)

    async def run_the_update_ytd(self, **params: Any) -> None:
        await self.run("ytd", **params)

    async def run_the_update_itd(self, **params: Any) -> None:
        await self.run("itd", **params)

    async def run_the_update_range(
        self, *, start_date: str, end_date: str, **params: Any
    ) -> None:
        params = dict(params)
        params.update({"start_date": start_date, "end_date": end_date})
        await self.run("range", **params)

    # ---------------- Validation & Cleanup ----------------

    @staticmethod
    def _validate_range_params_if_needed(
        period: str, params: Mapping[str, Any]
    ) -> None:
        """Ensure required params exist for range runs."""
        if period != "range":
            return
        sd = params.get("start_date")
        ed = params.get("end_date")
        if not (sd and ed):
            raise ValueError(
                "For period='range' you must provide start_date and end_date (YYYY-MM-DD)."
            )
        try:
            dt.date.fromisoformat(sd)
            dt.date.fromisoformat(ed)
        except (ValueError, TypeError) as e:
            raise ValueError(f"Invalid start_date/end_date: {e}")

    # ---------------- Cleanup & Context Management ----------------

    def _cleanup(self) -> None:
        """Synchronous cleanup."""
        try:
            if self._orchestrator:
                self._orchestrator.close()  # ManagedResource.close() is sync
        except Exception:
            self.logger.warning("Error cleaning up orchestrator", exc_info=True)
        self._close_dask_client()

    async def _acleanup(self) -> None:
        """Asynchronous cleanup."""
        try:
            if self._orchestrator:
                # Assuming ManagedResource doesn't natively have aclose yet unless we added it,
                # but ArtifactOrchestrator (which inherits ManagedResource) might reuse sync cleanup
                # or we rely on ManagedResource standard cleanup.
                # Since ArtifactOrchestrator manages a Dask client connection sometimes, explicit cleanup is good.
                if hasattr(self._orchestrator, "aclose"):
                    await self._orchestrator.aclose()
                else:
                    self._orchestrator.close()
        except Exception:
            self.logger.warning("Error cleaning up orchestrator", exc_info=True)
        self._close_dask_client()

    async def __aenter__(self):
        """
        Async Context Entry.
        Start Dask cluster in a separate thread to avoid blocking the event loop.
        """
        if not self.dask_client and self.own_dask_client is False:
            import asyncio

            await asyncio.to_thread(self._start_dask_cluster, **self._lazy_init_kwargs)

        # ManagedResource.__aenter__ typically not needed unless it has async setup,
        # but good practice if available.
        # await super().__aenter__()
        return self

    async def __aexit__(self, exc_type, exc, tb) -> bool:
        """Ensure async cleanup always runs on exit."""
        await self.aclose()
        return False

    async def aclose(
        self, *, suppress_errors: bool = False, run_sync_cleanup_if_missing: bool = True
    ) -> None:
        """Explicit async shutdown."""
        try:
            await self._acleanup()
            # If ManagedResource has aclose, call it. As of now it's usually sync close.
            # self.close()
        except Exception as e:
            if not suppress_errors:
                raise
            self.logger.warning(
                "PipelineExecutor async close suppressed error", exc_info=True
            )
