"""
Parquet Executor module.

Manages threaded execution of data loading and saving tasks for backfill operations.
"""

from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Type, Dict, List, Callable
import time
import datetime
import threading
import functools
from tqdm import tqdm
import pandas as pd
from sibi_flux.core import SecureResource
from sibi_flux.parquet import ParquetSaver

from sibi_flux.utils import ensure_slash
from sibi_flux.utils.retry import with_retry


class DataWrapper(SecureResource):
    """
    Wrapper for dataset providers to handle batch processing and persistence.
    """

    logger_extra = {"sibi_flux_component": "warehouse.data_wrapper"}

    def __init__(self, provider_class: Type, date_field: str, data_path: str, **kwargs):
        super().__init__(**kwargs)
        self.provider_class = provider_class
        self.date_field = date_field
        self.data_path = ensure_slash(data_path)

        self.show_progress = kwargs.get("show_progress", False)
        self.timeout = kwargs.get("timeout", 30)
        self.max_threads = kwargs.get("max_threads", 3)

        self.class_params = kwargs.get("class_params") or {
            "debug": self.debug,
            "logger": self.logger,
            "fs": self.fs,
            "verbose": self.verbose,
        }
        self.load_params = kwargs.get("load_params") or {}

        self._lock = threading.Lock()
        self.processed_dates: List[datetime.date] = []
        self.benchmarks: Dict[datetime.date, Dict[str, float]] = {}
        self.mmanifest = kwargs.get("mmanifest")
        self.update_planner = kwargs.get("update_planner")

        self._stop_event = threading.Event()
        self.logger_extra.update({"provider_class": self.provider_class.__name__})

    def _cleanup(self):
        self._stop_event.set()

    def process(self, max_retries=3, backoff_base=2.0, **kwargs):
        if not self.update_planner:
            return

        tasks = list(self.update_planner.get_tasks_by_priority())
        if not tasks:
            self.logger.info("No updates required.", extra=self.logger_extra)
            return

        if self.show_progress:
            self.update_planner.show_update_plan()

        start_t = time.perf_counter()

        # Create a partial function for the worker to avoid passing args repeatedly
        worker_fn = functools.partial(
            self._process_date_wrapper,
            retries=max_retries,
            backoff=backoff_base,
            **kwargs,
        )

        try:
            for priority, dates in tasks:
                if self._stop_event.is_set():
                    break
                self._execute_batch(priority, dates, worker_fn)
        finally:
            duration = time.perf_counter() - start_t
            count = len(self.processed_dates)
            if count > 0:
                self.logger.info(
                    f"Processed {count} in {duration:.1f}s", extra=self.logger_extra
                )
                if self.show_progress:
                    self.show_benchmark_summary()

    def _execute_batch(
        self, priority: int, dates: List[datetime.date], worker_fn: Callable
    ):
        max_workers = min(len(dates), self.max_threads)

        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            future_to_date = {
                executor.submit(worker_fn, date): date
                for date in dates
                if not self._stop_event.is_set()
            }

            iterator = as_completed(future_to_date)
            if self.show_progress:
                iterator = tqdm(
                    iterator, total=len(dates), desc=f"Priority {priority}", leave=False
                )

            for future in iterator:
                date = future_to_date[future]
                try:
                    future.result(timeout=self.timeout)
                except Exception as e:
                    self.logger.error(
                        f"Permanent failure for {date}: {e}", extra=self.logger_extra
                    )

    def _process_date_wrapper(
        self, date: datetime.date, retries: int, backoff: float, **kwargs
    ):
        decorator = with_retry(retries, backoff, stop_event=self._stop_event)
        return decorator(self._process_single_date)(date)

    def _process_single_date(self, date: datetime.date):
        is_hive = getattr(self, "partition_on_date", True)

        if is_hive:
            path = self.data_path
            log_ctx = {**self.logger_extra, "date": str(date)}
        else:
            path = f"{self.data_path}{date.year}/{date.month:02d}/{date.day:02d}/"
            log_ctx = {**self.logger_extra, "date": str(date), "legacy_path": path}

        if (
            self.update_planner
            and path in self.update_planner.skipped
            and getattr(self.update_planner, "ignore_missing", False)
        ):
            return

        t0 = time.perf_counter()
        load_params = {**self.load_params, f"{self.date_field}__date": date.isoformat()}

        with self.provider_class(**self.class_params) as instance:
            df = instance.load(**load_params)
            count = getattr(
                instance, "total_records", len(df) if hasattr(df, "__len__") else -1
            )

        t1 = time.perf_counter()

        if count == 0:
            if self.mmanifest:
                self.mmanifest.record(path)
            self.logger.info(f"No data for {date}", extra=log_ctx)
            return

        save_params = {
            "df_result": df,
            "parquet_storage_path": path,
            "fs": self.fs,
            "logger": self.logger,
            "debug": self.debug,
        }

        if is_hive:
            df["partition_date"] = date.isoformat()
            save_params["partition_on"] = ["partition_date"]

        with ParquetSaver(**save_params) as ps:
            ps.save_to_parquet()

        t2 = time.perf_counter()

        self._record_success(date, t1 - t0, t2 - t1, path)

    def _record_success(self, date, load_time, save_time, path):
        with self._lock:
            self.processed_dates.append(date)
            self.benchmarks[date] = {
                "load": load_time,
                "save": save_time,
                "total": load_time + save_time,
            }
        self.logger.info(
            f"Completed {date} (Load: {load_time:.1f}s, Save: {save_time:.1f}s)",
            extra={**self.logger_extra, "date": str(date)},
        )

    def show_benchmark_summary(self):
        if not self.benchmarks:
            return
        try:
            df = pd.DataFrame.from_dict(self.benchmarks, orient="index")
            df.index.name = "date"
            df.sort_index(inplace=True)
            self.logger.info(f"\n{df.to_string()}")
        except Exception:
            pass
