"""export_results handler — write datasets/results to disk."""
import json
from pathlib import Path

import polars as pl

from ..session_store import SessionStore


async def _handle_export_results(
    result_id: str,
    output_path: str,
    format: str = 'csv',
    overwrite: bool = False,
) -> dict:
    """Export a dataset (or result) to disk."""
    store = SessionStore()

    out = Path(output_path)
    file_paths: list[str] = []

    # Check if it's a dataset
    dataset_entry = store.get_dataset(result_id)
    if dataset_entry is not None:
        lf: pl.LazyFrame = dataset_entry['lazyframe']
        df = lf.collect(streaming=True)
        if out.suffix == '':
            out = out / f"export_{result_id}.{format}"

        if out.exists() and not overwrite:
            return {'error': f'Output exists and overwrite=False: {out}',
                    'truncated': False}

        out.parent.mkdir(parents=True, exist_ok=True)

        if format == 'csv':
            df.write_csv(out)
        elif format == 'parquet':
            df.write_parquet(out)
        elif format == 'json':
            df.write_json(out)
        else:
            return {'error': f'Unsupported format: {format}',
                    'truncated': False}

        file_paths.append(str(out.resolve()))

        return {'file_paths': file_paths, 'format': format,
                'row_count': len(df)}

    # Check if it's a cluster result
    cluster_entry = store.get_cluster_result(result_id)
    if cluster_entry is not None:
        if out.suffix == '':
            out = out / f"cluster_{result_id}.json"
        out.parent.mkdir(parents=True, exist_ok=True)

        if out.exists() and not overwrite:
            return {'error': f'Output exists and overwrite=False: {out}',
                    'truncated': False}

        export_data = {
            'result_id': result_id,
            'n_clusters': cluster_entry['n_clusters'],
            'params': cluster_entry['params'],
            'labels_sample': cluster_entry['labels'][:1000],
        }
        with open(out, 'w', encoding='utf-8') as f:
            json.dump(export_data, f, default=str)
        file_paths.append(str(out.resolve()))

        return {'file_paths': file_paths, 'format': 'json'}

    # Check feature result
    feat_entry = store.get_feature_result(result_id)
    if feat_entry is not None:
        if out.suffix == '':
            out = out / f"features_{result_id}.json"
        out.parent.mkdir(parents=True, exist_ok=True)

        if out.exists() and not overwrite:
            return {'error': f'Output exists and overwrite=False: {out}',
                    'truncated': False}

        with open(out, 'w', encoding='utf-8') as f:
            json.dump(feat_entry['features'], f, default=str)
        file_paths.append(str(out.resolve()))
        return {'file_paths': file_paths, 'format': 'json'}

    return {'error': f'No dataset or result found for ID: {result_id}',
            'truncated': False}
