"""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 _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).
    """
    appdata = os.environ.get('APPDATA', os.path.expanduser('~'))
    store_dir = Path(appdata) / 'TianXuan'
    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()
        return cls._instance

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

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

        Only metadata (schema, csv_glob, row_count, …) is saved — never the
        LazyFrame itself, which is an in-memory query plan.
        """
        with self._datalock:
            serializable = {}
            for ds_id, entry in self._stores.items():
                if entry.get('type') != 'dataset':
                    continue
                meta = entry.get('metadata', {})
                serializable[ds_id] = {
                    'schema': entry.get('schema'),
                    'metadata': 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_from_disk(self) -> int:
        """Re-register datasets from persisted metadata.

        Attempts two strategies in order:
        1. **SQLite table** — when ``metadata.sqlite_table`` exists (fast,
           always available, no CSV dependency).
        2. **CSV glob** — falls back to ``load_csv_directory`` for datasets
           that were stored before the SQLite migration or whose SQLite table
           was dropped.

        Returns the number of successfully restored datasets.
        """
        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

        # Lazy imports to avoid circular dependency at module level
        from analysis.data_loader import load_csv_directory, load_from_db  # fmt: skip

        restored = 0
        for ds_id, info in data.items():
            meta = dict(info.get('metadata', {}))
            schema = dict(info.get('schema', {}))
            restored_ok = False

            # Strategy 1: SQLite table (fast path)
            sqlite_table = meta.get('sqlite_table')
            if sqlite_table:
                try:
                    lf = load_from_db(sqlite_table)
                    if lf is not None:
                        self.store_dataset(ds_id, lf, schema=schema, metadata=meta)
                        restored += 1
                        restored_ok = True
                except Exception:
                    pass  # fall through to CSV restore

            if restored_ok:
                continue

            # Strategy 2: CSV glob (legacy / fallback)
            csv_glob = info.get('csv_glob')
            if not csv_glob:
                continue
            try:
                lf, schema_csv, row_count, file_count, memory_mb = load_csv_directory(csv_glob)
                # Merge schema from CSV (may be more accurate than persisted schema)
                self.store_dataset(ds_id, lf, schema=schema_csv, metadata=meta)
                restored += 1
            except Exception:
                # File(s) no longer exist or schema changed — skip gracefully
                continue
        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*)."""
        with self._datalock:
            entry = self._stores.get(dataset_id)
            if entry and entry.get('type') == 'dataset':
                return entry
            return None

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

        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]
        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'),
                    'columns': list(entry.get('schema', {}).keys()),
                    'column_count': len(entry.get('schema', {})),
                })
        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)}>"