"""Shared utilities for tool handlers.

Provides:
- NUMERIC_DTYPES / NUMERIC_TYPE_NAMES: shared constants for dtype filtering
- _json_size, _truncate_response: LLM context-window helpers
- _build_filter_expr: Polars filter expression builder
- _resolve_dataset: one-liner SessionStore lookup (replaces 30+ repeated blocks)
- _count_rows: fast streaming row count
- _get_numeric_columns: extract numeric column names from schema
"""
import json
import logging
import traceback
from typing import Any

import polars as pl

from ..session_store import SessionStore
logger = logging.getLogger(__name__)


# ═══════════════════════════════════════════════════════════════════════
#  Shared constants
# ═══════════════════════════════════════════════════════════════════════

NUMERIC_DTYPES = (
    pl.Int8, pl.Int16, pl.Int32, pl.Int64,
    pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64,
    pl.Float32, pl.Float64,
)

NUMERIC_TYPE_NAMES = {
    'Int8', 'Int16', 'Int32', 'Int64',
    'UInt8', 'UInt16', 'UInt32', 'UInt64',
    'Float32', 'Float64',
}


# ═══════════════════════════════════════════════════════════════════════
#  LLM response helpers
# ═══════════════════════════════════════════════════════════════════════

_MAX_RESPONSE_KB = 64


def _json_size(obj: Any) -> int:
    """Estimate JSON byte size of *obj*."""
    return len(json.dumps(obj, default=str, ensure_ascii=False))


def _truncate_response(result: dict, max_kb: int = _MAX_RESPONSE_KB) -> dict:
    """Truncate large lists/keys in *result* to stay under *max_kb*."""
    if _json_size(result) <= max_kb * 1024:
        result.setdefault('truncated', False)
        return result

    # Drop verbose keys first
    for costly_key in ('column_stats', 'features', 'correlation_matrix',
                        'labels', 'cluster_labels'):
        if costly_key in result and isinstance(result[costly_key], dict):
            result[costly_key] = {'__truncated__': True,
                                  'key_count': len(result[costly_key])}
        elif costly_key in result and isinstance(result[costly_key], list):
            lst = result[costly_key]
            if len(lst) > 10:
                result[costly_key] = lst[:10] + [{'__truncated__': True,
                                                   'omitted': len(lst) - 10}]
    result['truncated'] = True
    return result


# ═══════════════════════════════════════════════════════════════════════
#  Filter expression builder
# ═══════════════════════════════════════════════════════════════════════

def _build_filter_expr(column: str, op: str, value: Any) -> pl.Expr:
    """Build a single Polars expression from a filter descriptor."""
    col = pl.col(column)
    op_map = {
        'eq': col == value,
        'neq': col != value,
        'gt': col > value,
        'gte': col >= value,
        'lt': col < value,
        'lte': col <= value,
        'contains': col.str.contains(value) if value is not None else col.cast(str).str.contains(''),
        'not_contains': ~(col.str.contains(value) if value is not None else col.cast(str).str.contains('')),
        'in': col.is_in(value) if isinstance(value, (list, tuple)) else col.is_in([value]),
        'not_in': ~(col.is_in(value) if isinstance(value, (list, tuple)) else col.is_in([value])),
        'is_null': col.is_null(),
        'is_not_null': col.is_not_null(),
    }
    expr = op_map.get(op)
    if expr is None:
        raise ValueError(f"Unsupported filter operator: {op}")
    return expr


# ═══════════════════════════════════════════════════════════════════════
#  SessionStore helpers
# ═══════════════════════════════════════════════════════════════════════

def _resolve_dataset(dataset_id: str) -> tuple:
    """Resolve a dataset from SessionStore.

    Returns:
        (entry, None) on success, where *entry* is the SessionStore entry dict.
        (None, error_dict) if the dataset is not found.
    """
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    if entry is None:
        return None, {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
    return entry, None


def _count_rows(lf: pl.LazyFrame) -> int:
    """Count rows in a LazyFrame via fast streaming count."""
    try:
        return lf.select(pl.len()).collect(streaming=True).item()
    except Exception:
        logger.error("_count_rows failed: {}".format(traceback.format_exc()))
        return 0


def _get_numeric_columns(lf: pl.LazyFrame, schema: dict, max_cols: int = 10) -> list[str]:
    """Get numeric column names from a dataset.

    Prefers *schema* (string-keyed); falls back to LazyFrame collect_schema().
    """
    if schema:
        return [c for c in schema
                if schema[c].split('(')[0].strip() in NUMERIC_TYPE_NAMES
                and not c.startswith('_')][:max_cols]
    # Fallback to LazyFrame schema
    available = lf.collect_schema()
    return [available.names()[i] for i, dt in enumerate(available.dtypes())
            if dt in NUMERIC_DTYPES and not available.names()[i].startswith('_')][:max_cols]
