"""Diagnostic tool handlers — validate, explore, diagnose, compare, repair."""
import traceback
from analysis.constants import LOW_MEMORY_THRESHOLD_ROWS, RANDOM_SEED
import polars as pl

from ..session_store import SessionStore
from ._helpers import _resolve_dataset, _count_rows, NUMERIC_TYPE_NAMES
import logging
logger = logging.getLogger(__name__)


async def _handle_validate_data(dataset_id: str, fix_issues: bool = False) -> dict:
    """Check data quality: schema consistency, missing values, type conflicts."""
    entry, err = _resolve_dataset(dataset_id)
    if err:
        return err
    lf = entry['lazyframe']
    schema = entry.get('schema', {})
    issues = []
    df_sample = lf.head(1000).collect(streaming=True)
    for col in schema:
        null_count = df_sample[col].is_null().sum()
        if null_count > 500:
            issues.append(f'{col}: {null_count/10:.0f}% null')
        try:
            unique = df_sample[col].n_unique()
            if unique <= 1 and null_count == 0:
                issues.append(f'{col}: constant value (no variation)')
        except Exception:
            logger.error("_handle_validate_data failed: {}".format(traceback.format_exc()))
            pass
    total_rows = _count_rows(lf)
    return {
        'status': 'ok',
        'total_rows': total_rows,
        'columns': len(schema),
        'issues_found': len(issues),
        'issues': issues[:20],
        'recommendation': 'Try explore_distributions for detailed column analysis' if issues else 'No issues found',
    }


async def _handle_explore_distributions(dataset_id: str, columns: list[str] = None,
                                        max_sample: int = 10000) -> dict:
    """Per-column distribution statistics."""
    entry, err = _resolve_dataset(dataset_id)
    if err:
        return err
    lf = entry['lazyframe']
    schema = entry.get('schema', {})

    if not columns:
        columns = [c for c in schema if schema[c].split('(')[0].strip() in NUMERIC_TYPE_NAMES][:10]
    if len(columns) > 20:
        columns = columns[:20]

    try:
        n = _count_rows(lf)
        if n > max_sample:
            sample_lf = lf.select(columns).sample(n=max_sample, seed=RANDOM_SEED)
        else:
            sample_lf = lf.select(columns)
        df = sample_lf.collect(streaming=True)
    except Exception as exc:
        return {'error': f'explore_distributions failed: {exc}', 'truncated': False}

    stats = {}
    for col in columns:
        try:
            s = df[col]
            stats[col] = {
                'dtype': str(s.dtype),
                'null_pct': round(float(s.is_null().mean() * 100), 2),
                'unique': int(s.n_unique()),
                'top_values': [str(v) for v in s.head(5).to_list()],
            }
            if s.dtype in (pl.Float32, pl.Float64, pl.Int32, pl.Int64):
                stats[col]['min'] = float(s.min()) if s.min() is not None else None
                stats[col]['max'] = float(s.max()) if s.max() is not None else None
                stats[col]['mean'] = round(float(s.mean()), 4) if s.mean() is not None else None
                stats[col]['std'] = round(float(s.std()), 4) if s.std() is not None else None
        except Exception:
            logger.error("unknown failed: {}".format(traceback.format_exc()))
            stats[col] = {'error': 'failed to compute'}
    return {'columns': stats, 'total_rows': n, 'sampled': min(n, max_sample)}


async def _handle_find_outliers(dataset_id: str, columns: list[str] = None,
                                threshold: float = 3.0) -> dict:
    """Find statistical outliers using IQR method."""
    entry, err = _resolve_dataset(dataset_id)
    if err:
        return err
    lf = entry['lazyframe']
    schema = entry.get('schema', {})
    if not columns:
        columns = [c for c in schema if schema[c].split('(')[0].strip() in NUMERIC_TYPE_NAMES]
    n = _count_rows(lf)
    sample_lf = lf if n <= 100000 else lf.collect(streaming=True).sample(n=LOW_MEMORY_THRESHOLD_ROWS, seed=RANDOM_SEED).lazy()
    df = sample_lf.select(columns).collect(streaming=True)

    outlier_counts = {}
    for col in columns:
        try:
            s = df[col]
            q1, q3 = s.quantile(0.25), s.quantile(0.75)
            iqr = q3 - q1
            lo = q1 - threshold * iqr
            hi = q3 + threshold * iqr
            n_out = ((s < lo) | (s > hi)).sum()
            if n_out > 0:
                outlier_counts[col] = int(n_out)
        except Exception:
            logger.error("_handle_find_outliers failed: {}".format(traceback.format_exc()))
            pass
    return {'outlier_counts': outlier_counts, 'total_sampled': len(df), 'threshold_iqr': threshold}


async def _handle_diagnose_clustering(dataset_id: str, cluster_result_id: str) -> dict:
    """Diagnose clustering quality issues."""
    from sklearn.decomposition import TruncatedSVD
    store = SessionStore()
    entry = store.get_dataset(dataset_id)
    cluster_entry = store.get_cluster_result(cluster_result_id)
    if entry is None:
        return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
    if cluster_entry is None:
        return {'error': f'Cluster result not found: {cluster_result_id}', 'truncated': False}
    lf = entry['lazyframe']
    schema = entry.get('schema', {})
    num_cols = [c for c in schema if schema[c].split('(')[0].strip() in NUMERIC_TYPE_NAMES and not c.startswith('_')]

    n_clusters = cluster_entry.get('n_clusters', 0)
    n_noise = cluster_entry.get('n_noise', 0)
    quality = cluster_entry.get('quality_metrics', {})
    labels = cluster_entry.get('labels', [])
    params = cluster_entry.get('params', {})
    algo = params.get('algorithm', 'unknown')
    feature_cols = params.get('feature_columns', [])

    diagnosis = []
    if n_clusters <= 1:
        diagnosis.append('Only 1 cluster — data may lack structure or need feature scaling')
    if n_noise > 0.5 * len(labels) if labels else False:
        diagnosis.append(f'High noise ratio ({n_noise}/{len(labels)}) — try HDBSCAN with larger min_cluster_size')
    sil = quality.get('silhouette_score')
    if sil and sil < 0.2:
        diagnosis.append(f'Low silhouette ({sil:.3f}) — clusters overlap, try KMeans or feature selection')
    if len(feature_cols) < 3:
        diagnosis.append(f'Only {len(feature_cols)} features — add more columns for meaningful clusters')

    try:
        df_sample = lf.select(num_cols).head(5000).collect(streaming=True)
        mat = df_sample.to_numpy()
        n_comp = min(mat.shape)
        svd = TruncatedSVD(n_components=n_comp, random_state=RANDOM_SEED).fit(mat)
        var_explained = svd.explained_variance_ratio_
        if var_explained[0] > 0.95:
            diagnosis.append(f'First SVD component explains {var_explained[0]:.0%} variance — data nearly 1D')
    except Exception:
        logger.error("unknown failed: {}".format(traceback.format_exc()))
        pass

    return {
        'diagnosis': diagnosis,
        'n_clusters': n_clusters,
        'n_noise': n_noise,
        'silhouette': sil,
        'algorithm': algo,
        'n_features': len(feature_cols),
        'suggestions': [
            f'Try algorithm={"kmeans" if algo == "hdbscan" else "hdbscan"}',
            'Increase min_cluster_size for fewer, larger clusters',
            'Use explore_distributions to check feature quality',
        ],
    }


async def _handle_compare_datasets(dataset_id_a: str, dataset_id_b: str) -> dict:
    """Compare two datasets."""
    store = SessionStore()
    entry_a, entry_b = store.get_dataset(dataset_id_a), store.get_dataset(dataset_id_b)
    if entry_a is None or entry_b is None:
        return {'error': 'One or both datasets not found', 'truncated': False}

    def _get_info(entry):
        schema = entry.get('schema', {})
        lf = entry['lazyframe']
        try:
            n = _count_rows(lf)
        except Exception:
            logger.error("_get_info failed: {}".format(traceback.format_exc()))
            n = -1
        return {'schema': set(schema.keys()), 'row_count': n, 'columns': len(schema)}

    info_a, info_b = _get_info(entry_a), _get_info(entry_b)
    common_cols = info_a['schema'] & info_b['schema']
    only_a = info_a['schema'] - info_b['schema']
    only_b = info_b['schema'] - info_a['schema']
    return {
        'dataset_a': {'id': dataset_id_a, **info_a},
        'dataset_b': {'id': dataset_id_b, **info_b},
        'common_columns': len(common_cols),
        'columns_only_in_a': list(only_a)[:10] if only_a else [],
        'columns_only_in_b': list(only_b)[:10] if only_b else [],
        'row_diff': info_a['row_count'] - info_b['row_count'],
        'identical_schema': len(only_a) == 0 and len(only_b) == 0,
    }


async def _handle_export_debug_sample(dataset_id: str, max_rows: int = 50) -> dict:
    """Export raw data sample as JSON for debugging."""
    entry, err = _resolve_dataset(dataset_id)
    if err:
        return err
    df = entry['lazyframe'].head(max_rows).collect(streaming=True)
    rows = df.to_dict(as_series=False)
    result = []
    for i in range(len(df)):
        row = {}
        for col in df.columns:
            val = rows[col][i]
            if isinstance(val, (float, int, str, bool)):
                row[col] = val
            elif val is None:
                row[col] = None
            else:
                row[col] = str(val)
        result.append(row)
    return {'rows': result, 'count': len(result), 'columns': df.columns}


async def _handle_repair_schema(dataset_id: str) -> dict:
    """Attempt to repair schema mismatches by normalising column names."""
    entry, err = _resolve_dataset(dataset_id)
    if err:
        return err
    lf = entry['lazyframe']
    schema = entry.get('schema', {})
    col_map = {}
    for c in schema:
        norm = c.lower().replace('-', '_').replace('.', '_')
        col_map[c] = norm
    if len(set(col_map.values())) < len(col_map):
        return {'status': 'conflict', 'message': 'Normalised names conflict', 'column_map': col_map}
    renames = {old: new for old, new in col_map.items() if old != new}
    if renames:
        lf = lf.rename(renames)
        new_schema = {col_map.get(c, c): dtype for c, dtype in schema.items()}
        store = SessionStore()
        store.store_dataset(dataset_id, lf, schema=new_schema, metadata=entry.get('metadata', {}))
        return {'status': 'repaired', 'renamed': len(renames), 'column_map': col_map}
    return {'status': 'no_change', 'message': 'Column names already normalised'}
