"""Bluesky catalog implementation."""

from typing import Any, List, Optional, Dict
from qtpy.QtCore import QObject
from databroker.queries import TimeRange

from .base import CatalogBase
from ..data import BlueskyRun, NBSRun
from .chunkCache import ChunkCache
from nbs_viewer.utils import print_debug
from .worker_pool import CatalogWorkerPool


class BlueskyCatalog(CatalogBase):
    """
    Catalog implementation for Bluesky data.

    This class manages a collection of runs from a Bluesky catalog,
    providing access to run data and metadata.
    """

    def __init__(self, catalog: Any, parent: Optional[QObject] = None):
        """
        Initialize the catalog.

        Parameters
        ----------
        catalog : Any
            The Bluesky catalog object.
        parent : QObject, optional
            Parent object, by default None.
        """
        super().__init__(parent)
        self._base_catalog = catalog.v2
        self._catalog = catalog.v2
        self._wrapped_runs = {}
        self._chunk_cache = ChunkCache()
        # Enable async key discovery by default for Bluesky-based catalogs
        try:
            self._worker_pool: CatalogWorkerPool | None = CatalogWorkerPool(self)
        except Exception:
            self._worker_pool = None

    def __len__(self):
        return len(self._catalog)

    def __getattr__(self, name):
        return getattr(self._catalog, name)

    @property
    def columns(self) -> List[str]:
        """Get column names for display."""
        return BlueskyRun.to_header()

    def get_runs(self) -> List[BlueskyRun]:
        """Get all available runs."""
        runs = []
        for uid, run in self._catalog.items():
            if uid not in self._wrapped_runs:
                print_debug(
                    "BlueskyCatalog.get_runs",
                    f"Wrapping run {uid}",
                    "run",
                )
                self._wrapped_runs[uid] = self.wrap_run(run, uid)
            runs.append(self._wrapped_runs[uid])
        return runs

    def wrap_run(self, run, uid):
        # Reuse existing wrapper
        print_debug("BlueskyCatalog.wrap_run", f"Wrapping run {uid}", "run")

        br = BlueskyRun(run, uid, self._catalog, chunk_cache=self._chunk_cache)
        self._wrapped_runs[uid] = br
        print_debug(
            "BlueskyCatalog.wrap_run",
            f"wrapped_runs has len {len(self._wrapped_runs)}",
            "run",
        )
        # If a worker pool exists, start async key initialization once
        if self._worker_pool is not None and not getattr(
            br, "_keys_init_started", False
        ):
            br.keys_loading.emit()
            self._worker_pool.submit_key_init(br)
        elif self._worker_pool is None and not getattr(br, "_keys_initialized", False):
            # Fallback to synchronous initialization if no pool and not done yet
            br._initialize_keys()
        return br

    def get_run(self, uid: str) -> BlueskyRun:
        """
        Get a specific run by UID.

        Parameters
        ----------
        uid : str
            The run's unique identifier.

        Returns
        -------
        BlueskyRun
            The requested run.

        Raises
        ------
        KeyError
            If run not found.
        """
        if uid not in self._wrapped_runs:
            print_debug(
                "BlueskyCatalog.get_run",
                f"Wrapping run {uid}",
                "run",
            )
            print_debug(
                "BlueskyCatalog.get_run",
                f"wrapped_runs has len {len(self._wrapped_runs)}",
                "run",
            )
            if uid not in self._catalog:
                raise KeyError(f"No run found with UID {uid}")
            self._wrapped_runs[uid] = self.wrap_run(self._catalog[uid], uid)
        print_debug(
            "BlueskyCatalog.get_run",
            f"returning wrapped run for {uid}",
            "run",
        )
        return self._wrapped_runs[uid]

    def items_slice(self, slice_obj: Optional[slice] = None):
        """
        Returns a generator of (key, RUN_WRAPPER) pairs for a slice of items in the catalog.

        Parameters
        ----------
        slice_obj : slice, optional
            A slice object to apply to the catalog items. If None, returns all items.

        Yields
        ------
        Tuple[Any, RUN_WRAPPER]
            Key-value pairs where the value is wrapped in a RUN_WRAPPER.
        """
        print_debug(
            "BlueskyCatalog.items_slice",
            f"Getting slice {slice_obj}",
            category="DEBUG_RUNLIST",
        )
        sliced_items = (
            self._catalog.items()[slice_obj] if slice_obj else self._catalog.items()
        )
        for key, value in sliced_items:
            # print(f"Wrapping run {key}")
            try:
                if key not in self._wrapped_runs:
                    print_debug(
                        "BlueskyCatalog.items_slice",
                        f"Wrapping run {key}",
                        "run",
                    )
                    self._wrapped_runs[key] = self.wrap_run(value, key)
                yield key, self._wrapped_runs[key]
            except Exception as ex:
                print_debug(
                    "BlueskyCatalog.items_slice",
                    f"Error wrapping run {key}: {ex}",
                    category="DEBUG_RUNLIST",
                )

    def search(self, query: Dict) -> "BlueskyCatalog":
        """
        Search for runs matching query.

        Parameters
        ----------
        query : Dict
            Search criteria.

        Returns
        -------
        BlueskyCatalog
            New catalog containing matching runs.
        """
        self._catalog = self._base_catalog.search(query)
        return self

    def filter_by_time(
        self, since: Optional[str] = None, until: Optional[str] = None
    ) -> "BlueskyCatalog":
        """
        Filter runs by time range.

        Parameters
        ----------
        since : str, optional
            Start time (YYYY-MM-DD), by default None.
        until : str, optional
            End time (YYYY-MM-DD), by default None.

        Returns
        -------
        BlueskyCatalog
            New catalog with filtered runs.
        """
        return self.search(TimeRange(since=since, until=until))


class NBSCatalog(BlueskyCatalog):
    """Catalog implementation for NBS data."""

    def wrap_run(self, run, uid):
        print_debug("NBSCatalog.wrap_run", f"Wrapping run {uid}", "run")
        br = NBSRun(run, uid, self._catalog, chunk_cache=self._chunk_cache)
        print_debug("NBSRun", f"NBSRun for run {uid}: {id(br)}", "run")
        self._wrapped_runs[uid] = br
        print_debug(
            "NBSCatalog.wrap_run",
            f"wrapped_runs has len {len(self._wrapped_runs)}",
            "run",
        )
        if self._worker_pool is not None and not getattr(
            br, "_keys_init_started", False
        ):
            br.keys_loading.emit()
            self._worker_pool.submit_key_init(br)
        elif self._worker_pool is None and not getattr(br, "_keys_initialized", False):
            br._initialize_keys()
        return br

    @property
    def columns(self) -> List[str]:
        """Get column names for display."""
        return NBSRun.to_header()  # Use NBSRun headers instead of BlueskyRun
