"""compute_distance_matrix handler — LLM-driven per-row evaluation."""
import traceback
from analysis.constants import RANDOM_SEED
from typing import Optional

import numpy as np
import polars as pl

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


def _indent_function_body(body: str, indent: str = '    ') -> str:
    """Indent each non-empty line of a code block for wrapping in a def."""
    return '\n'.join(
        indent + line if line.strip() else line
        for line in body.split('\n')
    )


async def _handle_compute_distance_matrix(
    dataset_id: str,
    python_function: str,
    columns: Optional[list[str]] = None,
    max_sample: int = 10000,
) -> dict:
    """Evaluate a user-provided Python function on each row of a dataset."""
    entry, err = _resolve_dataset(dataset_id)
    if err:
        return err

    lf: pl.LazyFrame = entry['lazyframe']
    df = lf.collect(streaming=True)

    if columns is None:
        columns = [c for c in df.columns
                   if df[c].dtype in NUMERIC_DTYPES]

    available = set(df.columns)
    columns = [c for c in columns if c in available]
    if not columns:
        return {'error': 'No valid columns found in dataset',
                'truncated': False}

    n_total = len(df)

    if n_total > max_sample:
        df_sampled = df.sample(n=max_sample, seed=RANDOM_SEED)
        n_evaluated = max_sample
    else:
        df_sampled = df
        n_evaluated = n_total

    rows = df_sampled.select(columns).to_dicts()

    safe_builtins = {
        'abs': abs, 'len': len, 'range': range,
        'min': min, 'max': max, 'sum': sum, 'round': round,
        'int': int, 'float': float, 'str': str, 'bool': bool,
        'list': list, 'dict': dict, 'tuple': tuple,
        'isinstance': isinstance, 'type': type,
        'True': True, 'False': False, 'None': None,
        'enumerate': enumerate, 'zip': zip, 'map': map,
        'filter': filter, 'sorted': sorted, 'reversed': reversed,
    }

    safe_ns = {
        '__builtins__': safe_builtins,
        'math': __import__('math'),
        'np': __import__('numpy'),
    }

    body = python_function.strip()
    if body.startswith('def '):
        full_code = body
    else:
        indented = _indent_function_body(body)
        full_code = f'def distance_fn(row):\n{indented}'

    try:
        exec(full_code, safe_ns)
        distance_fn = safe_ns.get('distance_fn')
        if distance_fn is None:
            for _name, _obj in safe_ns.items():
                if callable(_obj) and _name not in ('math', 'np', '__builtins__'):
                    distance_fn = _obj
                    break
        if distance_fn is None:
            return {'error': 'Could not locate the compiled function',
                    'truncated': False}
    except SyntaxError as e:
        return {'error': f'Function syntax error: {e}', 'truncated': False}
    except Exception as e:
        return {'error': f'Function compilation failed: {e}',
                'truncated': False}

    scores: list[float | None] = []
    errors = 0
    for row in rows:
        try:
            score = distance_fn(row)
            if isinstance(score, (int, float)):
                scores.append(float(score))
            else:
                scores.append(None)
                errors += 1
        except Exception:
            logger.error("unknown failed: {}".format(traceback.format_exc()))
            scores.append(None)
            errors += 1

    valid_scores = [s for s in scores if s is not None]
    if not valid_scores:
        return {'error': 'No valid scores computed — all rows raised errors',
                'truncated': False}

    arr = np.array(valid_scores)
    summary = {
        'min': float(np.min(arr)),
        'max': float(np.max(arr)),
        'mean': float(np.mean(arr)),
        'std': float(np.std(arr)),
        'count': len(arr),
        'errors': errors,
        'total_rows_in_sample': len(rows),
        'matrix_shape': [len(rows), len(columns)],
    }

    scored_sample = []
    for i, row in enumerate(rows[:30]):
        if i < len(scores):
            scored_sample.append({
                'index': i,
                'values': {c: row.get(c) for c in columns[:3]},
                'score': scores[i],
            })

    return {
        'summary': summary,
        'scored_sample': scored_sample,
        'columns_used': columns,
        'total_rows_in_dataset': n_total,
        'rows_evaluated': n_evaluated,
    }
