from __future__ import absolute_import

__all__ = [
    'BaseEmbedder',
    'CategoryEmbedder',
    'NumericEmbedder',
    'SentenceTransformerEmbedder',
]

import os
import dbm.dumb as _dumbdbm
import shelve as she
import threading
import warnings
import logging
import numpy as np

_log = logging.getLogger(__name__)


def _open_shelf(filename):
    """Open a shelve using the dumb-dbm backend.

    Python's default shelve backend (gdbm on Linux) holds an exclusive file
    lock, which prevents two Jupyter kernels from sharing the same cache file.
    dbm.dumb uses plain text files and has no locking, making it safe for
    concurrent read/write access from multiple kernels.
    """
    return she.Shelf(
        dict=_dumbdbm.open(filename, 'c'),
        protocol=None,
        writeback=False)

# Default cache directory: <project_root>/.cache/pwml/
# Resolved from this file's location so it is always absolute and independent
# of the working directory (important when notebooks run from examples/).
# Structure: pwml/classifiers/embedders.py -> up 3 levels -> project root
# Override with the PWML_CACHE_DIR environment variable.
_DEFAULT_CACHE_DIR = (
    os.environ.get('PWML_CACHE_DIR')
    or os.path.join(
        os.path.dirname(  # project root
            os.path.dirname(  # pwml/
                os.path.dirname(  # pwml/classifiers/
                    os.path.abspath(__file__)))),
        '.cache', 'pwml'))

# Module-level registry of open shelve handles.  All BaseEmbedder instances
# that resolve to the same cache path share a single handle so that gdbm's
# exclusive file lock is never contested (e.g. when cross_validate creates
# multiple fold models in the same process, each with their own embedders).
_OPEN_SHELVES: dict = {}
_SHELVES_LOCK = threading.Lock()


class BaseEmbedder:
    """Base class for all embedding strategies.

    Thread safety
    -------------
    The shared shelve handle (``_OPEN_SHELVES``) is protected by
    ``_SHELVES_LOCK`` during :meth:`load_cache`.  However, individual calls to
    :meth:`embed_data` are **not** serialised.  Concurrent reads from the same
    shelve are safe (dbm.dumb is a plain text file), but concurrent writes from
    multiple threads to the same embedder instance may interleave.  If you use
    embedders in a multi-threaded context, create one embedder instance per
    thread or protect :meth:`embed_data` calls with an external lock.
    """

    def __init__(self, input_type=None, input_size=None, output_size=None, cache_path=None):

        self.input_type = input_type
        self.input_size = input_size
        self.output_size = output_size
        self.cache_path = cache_path
        self._shelve = None

        self.load_cache()

    def __enter__(self):
        return self

    def __exit__(self, *exc_info):
        self.close()

    def __del__(self):
        # Detach from the shared shelve handle but do not close it - the
        # handle is shared across instances and closing it here would corrupt
        # other live embedders that reference the same file.
        self._shelve = None

    def load_cache(self):
        # Cache files live in .cache/pwml/ at the project root (or in
        # cache_path if explicitly provided).  The directory is created on
        # demand so no manual setup is required.
        #
        # All embedder instances with the same resolved path share a single
        # shelve handle via _OPEN_SHELVES so that gdbm's exclusive file lock
        # is never contested (e.g. during cross_validate which creates many
        # fold models in the same process).
        base_dir = os.path.abspath(
            self.cache_path if self.cache_path is not None else _DEFAULT_CACHE_DIR)

        # If PWML_CACHE_ROOT is set, enforce that base_dir stays within it.
        # This guards against directory traversal attacks when cache_path comes
        # from an untrusted config file (e.g. "../../etc/").  When the env var
        # is not set no restriction is applied (paths are resolved to absolute
        # form above, so plain ".." escapes are normalised away).
        cache_root_env = os.environ.get('PWML_CACHE_ROOT')
        if cache_root_env:
            safe_root = os.path.abspath(cache_root_env)
            if not (base_dir.startswith(safe_root + os.sep) or base_dir == safe_root):
                raise ValueError(
                    "cache_path '{0}' is outside the allowed root '{1}' "
                    "(set by PWML_CACHE_ROOT).".format(base_dir, safe_root))

        os.makedirs(base_dir, exist_ok=True)
        self._cache_file = os.path.join(base_dir, '{0}.shelve'.format(self.input_type))
        with _SHELVES_LOCK:
            if self._cache_file not in _OPEN_SHELVES:
                _OPEN_SHELVES[self._cache_file] = _open_shelf(self._cache_file)
            self._shelve = _OPEN_SHELVES[self._cache_file]

    def close(self):
        """Detach this instance from the shared cache handle.

        The underlying shelve file stays open (shared with other instances).
        Call flush_cache() to explicitly sync it to disk.
        """
        self._shelve = None

    def flush_cache(self):
        """Sync the shared cache to disk without closing it."""
        if hasattr(self, '_cache_file') and self._cache_file in _OPEN_SHELVES:
            _OPEN_SHELVES[self._cache_file].sync()

    def data_empty(self, data, feature):
        raise NotImplementedError('Subclasses must implement data_empty()')

    def data_preparation(self, data, feature):
        return data

    def data_embedding(self, data, feature):
        return data

    def _cache_key(self, data, feature):
        # shelve requires string keys; include feature name to scope the cache
        # correctly when the same raw value can appear in different features
        # (e.g. the same string in two categorical columns with different classes).
        return '{0}::{1}'.format(feature.feature_name, str(data))

    def embed_data(self, data, feature):
        if data is None:
            return self.data_empty(
                data=data,
                feature=feature)

        data_prepared = self.data_preparation(
            data=data,
            feature=feature)

        cache_key = self._cache_key(data, feature)

        if cache_key in self._shelve:
            return self._shelve[cache_key]

        embedding = self.data_embedding(
            data=data_prepared,
            feature=feature)

        self._shelve[cache_key] = embedding

        return embedding


class CategoryEmbedder(BaseEmbedder):

    def __init__(self, input_size=None, output_size=None, cache_path=None):
        super().__init__(
            input_type='category',
            input_size=input_size,
            output_size=output_size,
            cache_path=cache_path)
        self._class_dicts: dict = {}

    def data_empty(self, data, feature):
        empty_data = None

        if self.output_size is None:
            empty_data = [0] * len(feature.classes)
        else:
            empty_data = [0] * self.output_size

        return np.array(
            empty_data,
            dtype=np.float64)

    def data_embedding(self, data, feature):
        vect = self.data_empty(data, feature)

        if data is not None:
            # Build a lookup dict once per feature (O(k) amortised over all samples).
            if feature.feature_name not in self._class_dicts:
                self._class_dicts[feature.feature_name] = {
                    v: i for i, v in enumerate(feature.classes)
                }
            position = self._class_dicts[feature.feature_name].get(data)
            if position is not None and position < len(vect):
                vect[position] = 1
            # Unknown category values produce a zero vector (all-absent encoding)
            # rather than raising an error, allowing inference on unseen data.

        return vect


_VALID_OOD_POLICIES = frozenset({'clip', 'warn', 'raise'})


class NumericEmbedder(BaseEmbedder):

    def __init__(self, default_value=0.0, cache_path=None, out_of_range='clip'):
        """Create a NumericEmbedder.

        Args:
            default_value (float): Value returned when the input is None.
            cache_path (str, optional): Directory for the shelve cache.
            out_of_range (str): Policy for inference-time values outside the
                training range.  One of:
                - ``'clip'`` (default): silently clamp to [0, 1] after normalisation.
                - ``'warn'``: normalise without clamping, emit a UserWarning.
                - ``'raise'``: raise a ValueError.
        """
        if out_of_range not in _VALID_OOD_POLICIES:
            raise ValueError(
                "out_of_range must be one of {0!r}, got {1!r}".format(
                    sorted(_VALID_OOD_POLICIES), out_of_range))
        super().__init__(
            input_type='numeric',
            input_size=None,
            output_size=None,
            cache_path=cache_path)

        self.default_value = default_value
        self.out_of_range = out_of_range

    def data_empty(self, data, feature):
        return np.array(
            [self.default_value],
            dtype=np.float64)

    def embed_data(self, data, feature):
        """Override to bypass the shelve cache.

        Numeric embedding is a single float operation - caching adds overhead
        with no benefit.  Normalisation also depends on training-time range, so
        cached raw values from a previous run would be incorrect.
        """
        if data is None:
            return self.data_empty(data=data, feature=feature)
        return self.data_embedding(data=data, feature=feature)

    def data_embedding(self, data, feature):
        """Embed a numeric value, normalising to [0, 1] using the training range.

        The range (min, max) is stored on the InputFeature during
        load_from_dataframe and serialised with the model, so it is available
        at inference time.  If the range is not set (e.g. a custom embedder
        setup), the raw value is used.

        Out-of-range values (below training min or above training max) are
        handled according to the ``out_of_range`` policy set at construction:
        ``'clip'`` clamps to [0, 1], ``'warn'`` normalises and emits a
        UserWarning, ``'raise'`` raises a ValueError.
        """
        val = float(data)
        if hasattr(feature, 'range') and feature.range is not None:
            min_val, max_val = feature.range
            if max_val > min_val:
                normalised = (val - min_val) / (max_val - min_val)
                if normalised < 0.0 or normalised > 1.0:
                    if self.out_of_range == 'raise':
                        raise ValueError(
                            "Feature '{0}': value {1} is outside the training "
                            "range [{2}, {3}].".format(
                                feature.feature_name, val, min_val, max_val))
                    elif self.out_of_range == 'warn':
                        warnings.warn(
                            "Feature '{0}': value {1} is outside the training "
                            "range [{2}, {3}]; normalised value {4:.4f} is "
                            "outside [0, 1].".format(
                                feature.feature_name, val, min_val, max_val,
                                normalised),
                            UserWarning,
                            stacklevel=3)
                    else:  # 'clip'
                        normalised = float(np.clip(normalised, 0.0, 1.0))
                val = normalised
        return np.array([val], dtype=np.float64)


# Known embedding dimensions for common sentence-transformers models.
_ST_KNOWN_DIMS = {
    'all-MiniLM-L6-v2':              384,
    'all-MiniLM-L12-v2':             384,
    'all-mpnet-base-v2':             768,
    'multi-qa-MiniLM-L6-cos-v1':     384,
    'multi-qa-mpnet-base-dot-v1':    768,
    'paraphrase-MiniLM-L6-v2':       384,
    'paraphrase-multilingual-MiniLM-L12-v2': 384,
    'all-distilroberta-v1':          768,
}


class SentenceTransformerEmbedder(BaseEmbedder):
    """Text embedder backed by sentence-transformers.

    Replaces the former TF Hub NNLM / Universal Sentence Encoder embedders.
    Works for both short product-description text ('text') and longer passages
    ('longtext') — the underlying transformer handles variable length natively.

    Args:
        input_type (str): Feature type key used for embedder registry lookup.
            Use 'text' for short descriptions, 'longtext' for paragraphs.
            Default 'text'.
        model_name (str): Any sentence-transformers model identifier.
            Default 'all-MiniLM-L6-v2' (384-dim, ~80 MB, fast on CPU).
            Use 'all-mpnet-base-v2' (768-dim, ~420 MB) for higher quality.
        cache_path (str, optional): Directory for the shelve embedding cache.

    Requirements:
        pip install sentence-transformers
    """

    def __init__(self, input_type='text', model_name='all-MiniLM-L6-v2', cache_path=None):
        self.model_name = model_name
        self._st_model = None  # lazy-loaded on first embed call

        output_size = _ST_KNOWN_DIMS.get(model_name, 384)

        super().__init__(
            input_type=input_type,
            input_size=None,
            output_size=output_size,
            cache_path=cache_path)

    def _get_model(self):
        """Lazy-load the sentence-transformers model (download on first call)."""
        if self._st_model is None:
            from sentence_transformers import SentenceTransformer
            self._st_model = SentenceTransformer(self.model_name)
            # Sync output_size with the actual model dimension.
            self.output_size = self._st_model.get_sentence_embedding_dimension()
        return self._st_model

    def data_empty(self, data, feature):
        return np.zeros(self.output_size, dtype=np.float64)

    def data_embedding(self, data, feature):
        model = self._get_model()
        embedding = model.encode(
            [str(data)],
            normalize_embeddings=True,
            show_progress_bar=False)
        return embedding[0].astype(np.float64)
