"""Thread-safe in-memory session store for datasets and results.

Singleton pattern with RLock guard. Stores:
  - Datasets: dataset_id -> (LazyFrame, schema, metadata)
  - Cluster results: result_id -> (labels, n_clusters, params, parent_dataset_id)
  - Feature results: result_id -> (features, parent_cluster_result_id)
"""
import json
import os
import threading
import gc
import uuid
from pathlib import Path
from typing import Any, Optional

import polars as pl


# Approximate per-element byte sizes for memory estimation
_DTYPE_BYTE_SIZES = {
    # Polars dtype strings
    'Int8': 1, 'Int16': 2, 'Int32': 4, 'Int64': 8,
    'UInt8': 1, 'UInt16': 2, 'UInt32': 4, 'UInt64': 8,
    'Float32': 4, 'Float64': 8,
    'Boolean': 1, 'Bool': 1,
    'Utf8': 50, 'String': 50, 'Categorical': 8,
    'Date': 4, 'Datetime': 8, 'Time': 8, 'Duration': 8,
    'List': 16, 'Object': 16, 'Null': 0,
    # Polars __str__ representations
    'i32': 4, 'i64': 8, 'u32': 4, 'u64': 8,
    'f32': 4, 'f64': 8, 'str': 50, 'bool': 1,
}


def _estimate_dtype_bytes(dtype: Any) -> int:
    """Guess per-element byte size from a Polars dtype."""
    key = str(dtype)
    for prefix, size in _DTYPE_BYTE_SIZES.items():
        if key.startswith(prefix) or key == prefix:
            return size
    return 8  # fallback


def _generate_id(prefix: str = 'ds') -> str:
    """Generate a short unique identifier."""
    return f"{prefix}_{uuid.uuid4().hex[:12]}"


def _estimate_row_count(lf: pl.LazyFrame) -> int:
    """Estimate row count from a LazyFrame (fast, streaming).

    Uses ``select(pl.len()).collect(streaming=True)`` which scans the
    underlying source efficiently without loading all columns.
    """
    try:
        return lf.select(pl.len()).collect(streaming=True).item()
    except Exception:
        return 0


def _get_store_path() -> Path:
    """Return the persistent JSON file path for dataset metadata.

    Uses ``%APPDATA%/TianXuan/.session_store.json`` (Windows) or
    ``~/.tianxuan_session_store.json`` (other platforms).
    """
    from analysis.appdata import get_appdata_dir
    store_dir = get_appdata_dir()
    store_dir.mkdir(parents=True, exist_ok=True)
    return store_dir / '.session_store.json'


class SessionStore:
    """Singleton in-memory store for datasets and analysis results.

    Thread-safe via RLock. Datasets store Polars LazyFrames (query plans,
    not materialised data), so clone/shallow-copy is effectively free.

    Dataset metadata is persisted to ``_get_store_path()`` so that after a
    server restart datasets whose CSV files still exist can be re-loaded.
    """

    _instance: Optional['SessionStore'] = None
    _lock: threading.RLock = threading.RLock()

    def __new__(cls) -> 'SessionStore':
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
                    cls._instance._stores: dict[str, dict[str, Any]] = {}
                    cls._instance._datalock = threading.RLock()
                    cls._instance._restore_done = threading.Event()
        return cls._instance

    # ── persistence ────────────────────────────────────────────────────

    def _save_to_disk(self) -> None:
        """Persist current dataset metadata to JSON file.

        Only metadata (schema, row_count, file_count, sqlite_table, …) is
        saved — never the LazyFrame itself.  The resulting file is typically
        under 10 KB even with many datasets.
        """
        with self._datalock:
            serializable = {}
            for ds_id, entry in self._stores.items():
                if entry.get('type') != 'dataset':
                    continue
                meta = entry.get('metadata', {})
                # Store only restore-critical fields; omit runtime noise
                slim_meta = {
                    k: v for k, v in meta.items()
                    if k in ('row_count', 'file_count', 'sqlite_table',
                              'csv_glob', 'columns', 'memory_mb',
                              'schema_overrides', 'svd_components')
                }
                serializable[ds_id] = {
                    'schema': entry.get('schema'),
                    'metadata': slim_meta,
                    'csv_glob': meta.get('csv_glob'),
                }
        path = _get_store_path()
        path.write_text(
            json.dumps(serializable, ensure_ascii=False, indent=2),
            encoding='utf-8',
        )

    def _restore_metadata(self) -> int:
        """Read metadata from disk and store as lightweight *stubs*.

        A stub contains everything needed to reconstruct the LazyFrame on
        demand but performs **no I/O** beyond reading the JSON file.  Real
        data loading is deferred until the first :meth:`get_dataset` call
        (or the background thread).

        Returns the number of stubs created.
        """
        path = _get_store_path()
        if not path.exists():
            return 0
        try:
            data = json.loads(path.read_text(encoding='utf-8'))
        except (json.JSONDecodeError, OSError):
            return 0

        count = 0
        with self._datalock:
            for ds_id, info in data.items():
                # Skip if already loaded (idempotent on re-calls)
                if ds_id in self._stores:
                    continue
                meta = dict(info.get('metadata', {}))
                schema = dict(info.get('schema', {}))
                csv_glob = info.get('csv_glob')
                self._stores[ds_id] = {
                    'type': 'dataset',
                    '_stub': {
                        'sqlite_table': meta.get('sqlite_table'),
                        'csv_glob': csv_glob,
                        'schema_overrides': meta.get('schema_overrides'),
                    },
                    'schema': schema,
                    'metadata': meta,
                    'parent_id': None,
                }
                count += 1
        self._restore_done.set()
        return count

    def _load_pending_dataset(self, ds_id: str) -> bool:
        """Resolve a stub entry into a real LazyFrame.

        Called on the first :meth:`get_dataset` for a dataset that was
        restored via :meth:`_restore_metadata`.

        Attempts two strategies:
        1.  **SQLite table** (fast — uses :func:`load_from_db_lazy`).
        2.  **CSV glob** (legacy fallback — scans the filesystem).

        Returns ``True`` if the dataset was successfully loaded.
        """
        from analysis.data_loader import load_csv_directory, load_from_db_lazy  # fmt: skip

        entry = self._stores.get(ds_id)
        if not entry or '_stub' not in entry:
            return True  # nothing to do

        stub = entry.pop('_stub')
        meta = entry.get('metadata', {})
        schema = entry.get('schema', {})

        # Strategy 1: SQLite table (fast, native Polars reader)
        sqlite_table = stub.get('sqlite_table')
        if sqlite_table:
            try:
                lf = load_from_db_lazy(sqlite_table, schema_overrides=stub.get('schema_overrides'))
                if lf is not None:
                    entry['lazyframe'] = lf
                    if 'row_count' not in meta:
                        meta['row_count'] = _estimate_row_count(lf)
                    return True
            except Exception:
                pass  # fall through to CSV restore

        # Strategy 2: CSV glob (legacy fallback)
        csv_glob = stub.get('csv_glob')
        if not csv_glob:
            # No recoverable data source — entry remains without LazyFrame
            return False
        try:
            lf, schema_csv, row_count, file_count, memory_mb = load_csv_directory(csv_glob)
            if 'file_count' not in meta and file_count is not None:
                meta['file_count'] = file_count
            if 'row_count' not in meta and row_count is not None:
                meta['row_count'] = row_count
            # Merge CSV-scanned schema (may have more/different columns)
            merged_schema = {**schema, **schema_csv}
            entry['lazyframe'] = lf
            entry['schema'] = merged_schema
            return True
        except Exception:
            return False

    def start_async_restore(self) -> int:
        """Restore dataset metadata immediately; load data in background.

        At startup this reads the small metadata JSON and stores *stubs*
        (metadata only, no data I/O).  A daemon thread then resolves each
        stub into a real LazyFrame without blocking the main thread.

        Returns the number of datasets queued for restoration.
        """
        count = self._restore_metadata()
        if count > 0:
            t = threading.Thread(
                target=self._background_restore,
                daemon=True,
                name='session-restore',
            )
            t.start()
        return count

    def _background_restore(self) -> None:
        """Pre-load all stubs in a daemon thread (best-effort)."""
        with self._datalock:
            ds_ids = [
                ds_id for ds_id, entry in self._stores.items()
                if entry.get('type') == 'dataset' and '_stub' in entry
            ]
        for ds_id in ds_ids:
            try:
                self._load_pending_dataset(ds_id)
            except Exception:
                pass  # best-effort; stub remains for on-demand fallback

    def restore_from_disk(self) -> int:
        """Legacy synchronous restore (kept for backward compatibility).

        Prefer :meth:`start_async_restore` for non-blocking startup.
        """
        self._restore_metadata()
        with self._datalock:
            ds_ids = [
                ds_id for ds_id, entry in self._stores.items()
                if entry.get('type') == 'dataset' and '_stub' in entry
            ]
        restored = 0
        for ds_id in ds_ids:
            if self._load_pending_dataset(ds_id):
                restored += 1
        return restored

    # ── dataset operations ──────────────────────────────────────────────

    def store_dataset(
        self,
        dataset_id: str,
        lazyframe: pl.LazyFrame,
        schema: Optional[dict] = None,
        metadata: Optional[dict] = None,
        parent_id: Optional[str] = None,
    ) -> None:
        """Store a LazyFrame under *dataset_id*.

        Args:
            dataset_id: Unique key for the dataset.
            lazyframe: Polars LazyFrame (query plan, not materialised).
            schema: Column schema dict {name: dtype_string}.
            metadata: Optional dict (row_count, file_count, …).
            parent_id: If this is a derived dataset (filter/clone/…), the
                source dataset_id.
        """
        with self._datalock:
            self._stores[dataset_id] = {
                'type': 'dataset',
                'lazyframe': lazyframe,
                'schema': schema or {},
                'metadata': metadata or {},
                'parent_id': parent_id,
            }
        self._save_to_disk()

    def get_dataset(self, dataset_id: str) -> Optional[dict]:
        """Retrieve a stored dataset entry (or *None*).

        If the entry is a *stub* (restored from metadata but not yet
        materialised), this method triggers on-demand loading of the
        LazyFrame before returning.
        """
        with self._datalock:
            entry = self._stores.get(dataset_id)
            if not entry or entry.get('type') != 'dataset':
                return None
            needs_load = '_stub' in entry

        if needs_load:
            self._load_pending_dataset(dataset_id)

        with self._datalock:
            return self._stores.get(dataset_id)

    def drop_dataset(self, dataset_id: str) -> bool:
        """Remove a dataset from the store and force garbage collection.

        Cascades to child datasets (e.g. filtered results) via parent_id tracking.
        Returns True if the dataset existed.
        """
        with self._datalock:
            existed = dataset_id in self._stores and self._stores[dataset_id].get('type') == 'dataset'
            if existed:
                del self._stores[dataset_id]
            # Cascade: drop all child datasets that reference the dropped one as parent
            child_ids = [
                k for k, v in self._stores.items()
                if v.get('type') == 'dataset' and v.get('parent_id') == dataset_id
            ]
            for cid in child_ids:
                del self._stores[cid]
        if existed:
            gc.collect()
            self._save_to_disk()
        return existed

    def clone_dataset(self, dataset_id: str) -> Optional[str]:
        """Shallow clone — LazyFrame is a query plan so no data is copied.

        Returns the new dataset_id, or *None* if source does not exist.
        """
        with self._datalock:
            entry = self.get_dataset(dataset_id)
            if entry is None:
                return None
            new_id = _generate_id('clone')
            self._stores[new_id] = {
                'type': 'dataset',
                'lazyframe': entry['lazyframe'],  # same plan pointer
                'schema': dict(entry['schema']),
                'metadata': dict(entry.get('metadata', {})),
                'parent_id': dataset_id,
            }
        return new_id

    def list_datasets(self) -> list[dict]:
        """Return summary dict for every stored dataset."""
        with self._datalock:
            results = []
            for ds_id, entry in self._stores.items():
                if entry.get('type') != 'dataset':
                    continue
                meta = entry.get('metadata', {})
                results.append({
                    'dataset_id': ds_id,
                    'parent_id': entry.get('parent_id'),
                    'row_count': meta.get('row_count'),
                    'file_count': meta.get('file_count'),
                    'svd_components': meta.get('svd_components', 0),
                    'columns': list(entry.get('schema', {}).keys()),
                    'column_count': len(entry.get('schema', {})),
                    'is_stub': '_stub' in entry,
                })
        return results

    def memory_estimate_mb(self, dataset_id: str) -> float:
        """Rough memory estimate based on schema × row_count."""
        with self._datalock:
            entry = self.get_dataset(dataset_id)
            if entry is None:
                return 0.0
            schema = entry.get('schema', {})
            row_count = entry.get('metadata', {}).get('row_count', 0)
        if not row_count or not schema:
            return 0.0
        bytes_per_row = sum(_estimate_dtype_bytes(dtype) for dtype in schema.values())
        return (bytes_per_row * row_count) / (1024 * 1024)

    # ── cluster-result operations ───────────────────────────────────────

    def store_cluster_result(
        self,
        result_id: str,
        labels: list[int],
        n_clusters: int,
        params: dict[str, Any],
        parent_dataset_id: str,
    ) -> None:
        """Store a clustering result."""
        with self._datalock:
            self._stores[result_id] = {
                'type': 'cluster_result',
                'labels': labels,
                'n_clusters': n_clusters,
                'params': params,
                'parent_dataset_id': parent_dataset_id,
            }

    def get_cluster_result(self, result_id: str) -> Optional[dict]:
        """Retrieve a stored cluster result."""
        with self._datalock:
            entry = self._stores.get(result_id)
            if entry and entry.get('type') == 'cluster_result':
                return entry
            return None

    def drop_cluster_result(self, result_id: str) -> bool:
        """Remove a cluster result."""
        with self._datalock:
            existed = result_id in self._stores and self._stores[result_id].get('type') == 'cluster_result'
            if existed:
                del self._stores[result_id]
        return existed

    # ── feature-result operations ───────────────────────────────────────

    def store_feature_result(
        self,
        result_id: str,
        features: list[dict[str, Any]],
        parent_cluster_result_id: str,
    ) -> None:
        """Store extracted features per cluster."""
        with self._datalock:
            self._stores[result_id] = {
                'type': 'feature_result',
                'features': features,
                'parent_cluster_result_id': parent_cluster_result_id,
            }

    def get_feature_result(self, result_id: str) -> Optional[dict]:
        """Retrieve a stored feature result."""
        with self._datalock:
            entry = self._stores.get(result_id)
            if entry and entry.get('type') == 'feature_result':
                return entry
            return None

    def drop_feature_result(self, result_id: str) -> bool:
        """Remove a feature result."""
        with self._datalock:
            existed = result_id in self._stores and self._stores[result_id].get('type') == 'feature_result'
            if existed:
                del self._stores[result_id]
        return existed

    # ── general ─────────────────────────────────────────────────────────

    def drop_all(self) -> int:
        """Drop every entry and run gc. Returns count of removed entries."""
        with self._datalock:
            count = len(self._stores)
            self._stores.clear()
        if count:
            gc.collect()
        return count

    def __repr__(self) -> str:
        with self._datalock:
            types = {}
            for entry in self._stores.values():
                t = entry.get('type', 'unknown')
                types[t] = types.get(t, 0) + 1
        parts = [f"{k}={v}" for k, v in sorted(types.items())]
        return f"<SessionStore {'; '.join(parts)}>"
