import threading
import uuid
from typing import Set

import fsspec
import pandas as pd


class MissingManifestManager:
    """
    Optimized manager for tracking missing data partitions.
    """

    def __init__(
        self,
        fs: fsspec.AbstractFileSystem,
        manifest_path: str,
        clear_existing: bool = False,
        **kwargs,
    ):
        self.fs = fs
        self.manifest_path = manifest_path
        # Use set for O(1) lookups during runtime, instead of constantly checking DF
        self._new_records: Set[str] = set()
        self._loaded_paths: Set[str] = set()
        self._loaded = False
        self._lock = threading.RLock()

        self.clear_existing = clear_existing
        self.logger = kwargs.get("logger")

        # Optimisation: Don't glob for orphans on init. It's slow on S3.

    def load_existing(self) -> Set[str]:
        """Lazy load the manifest."""
        with self._lock:
            if self._loaded:
                return self._loaded_paths

            if not self.clear_existing and self.fs.exists(self.manifest_path):
                try:
                    # Read only the 'path' column
                    df = pd.read_parquet(
                        self.manifest_path, columns=["path"], filesystem=self.fs
                    )
                    self._loaded_paths = set(df["path"].dropna().astype(str))
                except Exception as e:
                    if self.logger:
                        self.logger.warning(f"Failed to load manifest: {e}")

            self._loaded = True
            return self._loaded_paths

    def record(self, full_path: str) -> None:
        """Thread-safe recording."""
        if full_path:
            with self._lock:
                self._new_records.add(full_path)

    def save(self) -> None:
        """Atomic save merging new records with existing ones."""
        with self._lock:
            if not self._new_records and not self.clear_existing:
                return

            # Merge in memory first
            all_paths = self.load_existing() | self._new_records

            if not all_paths:
                return

            df = pd.DataFrame({"path": list(all_paths)})

            # Atomic Write Strategy
            temp_path = f"{self.manifest_path}.tmp-{uuid.uuid4().hex}"
            try:
                # Ensure dir exists (catch-all for different fs implementations)
                parent = self.manifest_path.rsplit("/", 1)[0]
                self.fs.makedirs(parent, exist_ok=True)

                df.to_parquet(temp_path, filesystem=self.fs, index=False)

                # Atomic rename if supported, else copy-delete
                try:
                    self.fs.rename(temp_path, self.manifest_path)
                except OSError:
                    self.fs.copy(temp_path, self.manifest_path)
                    self.fs.rm_file(temp_path)

                # Clear pending buffer
                self._new_records.clear()
                self.clear_existing = False

            except Exception as e:
                if self.logger:
                    self.logger.error(f"Manifest save failed: {e}")
                # Try cleanup temp
                try:
                    self.fs.rm_file(temp_path)
                except Exception:
                    pass
