"""
Data Routes - Data Upload and Query API Endpoints
==================================================
Flask blueprint for AnnData file upload and data retrieval endpoints.
"""

import os
import logging
import numpy as np
from flask import Blueprint, request, jsonify, send_file, make_response
from werkzeug.utils import secure_filename

from server.data_adaptor.anndata_adaptor import HighPerformanceAnndataAdaptor
from utils.adata_helpers import (
    canonical_embedding_keys as _canonical_embedding_keys,
    resolve_embedding_key as _resolve_embedding_key,
)
from utils.file_helpers import resolve_browse_path, is_h5ad_file


# Create blueprint
bp = Blueprint('data', __name__)


@bp.route('/memory', methods=['GET'])
def get_memory():
    """Return system and current-process memory usage for the memory bar."""
    try:
        import psutil, os as _os
        proc = psutil.Process(_os.getpid())
        process_mb = proc.memory_info().rss / (1024 * 1024)
        vm = psutil.virtual_memory()
        total_mb = vm.total / (1024 * 1024)
        used_mb  = vm.used  / (1024 * 1024)
        return jsonify({
            'process_mb': round(process_mb, 1),
            'total_mb':   round(total_mb, 1),
            'used_mb':    round(used_mb, 1),
        })
    except ImportError:
        # psutil not available – fall back to resource module (Unix)
        try:
            import resource
            rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
            # macOS returns bytes, Linux returns KB
            process_mb = rss / (1024 * 1024) if rss > 10**7 else rss / 1024
            return jsonify({'process_mb': round(process_mb, 1), 'total_mb': None, 'used_mb': None})
        except Exception:
            return jsonify({'process_mb': None, 'total_mb': None, 'used_mb': None})
    except Exception as e:
        logging.warning(f'Memory query failed: {e}')
        return jsonify({'process_mb': None, 'total_mb': None, 'used_mb': None})


@bp.route('/upload', methods=['POST'])
def upload_file():
    """Upload h5ad file and create data adaptor."""
    if 'file' not in request.files:
        return jsonify({'error': 'No file provided'}), 400

    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No file selected'}), 400

    if not file.filename.endswith('.h5ad'):
        return jsonify({'error': 'File must be .h5ad format'}), 400

    try:
        # Save uploaded file
        filename = secure_filename(file.filename)
        filepath = os.path.join(bp.upload_folder, filename)
        file.save(filepath)

        # Create high-performance data adaptor
        bp.state.current_adaptor = HighPerformanceAnndataAdaptor(filepath)
        # Expose underlying AnnData for tool endpoints
        bp.state.current_adata = bp.state.current_adaptor.adata
        bp.state.current_filename = filename
        # Normal (full) load — clear preview flag
        bp.state.is_preview_mode = False

        # Always store a deterministic random 2D embedding in obsm so it:
        #  - persists across tool runs (obsm is auto-subset on cell filtering)
        #  - survives page refresh (exposed via /api/status embeddings list)
        #  - is always available as a fallback before PCA/UMAP are computed
        if 'X_random' not in bp.state.current_adata.obsm:
            seed = abs(hash(filepath)) % (2**32 - 1)
            rng = np.random.RandomState(seed)
            bp.state.current_adata.obsm['X_random'] = rng.rand(
                bp.state.current_adata.n_obs, 2).astype(np.float32)

        try:
            bp.state.kernel_executor.sync_adata(bp.state.current_adata)
        except Exception:
            pass

        # Get schema and summary
        schema = bp.state.current_adaptor.get_schema()
        summary = bp.state.current_adaptor.get_data_summary()
        chunk_info = bp.state.current_adaptor.get_chunk_info()

        # Build response compatible with legacy UI expectations
        from utils.adata_helpers import analyze_data_state as _analyze_data_state
        adata = bp.state.current_adata
        embeddings = _canonical_embedding_keys(adata)
        response_data = {
            'filename': filename,
            # Legacy/UI fields used by single-cell.js
            'n_cells': adata.n_obs,
            'n_genes': adata.n_vars,
            'embeddings': embeddings,
            'obs_columns': list(adata.obs.columns),
            'var_columns': list(adata.var.columns),
            'uns_keys':    list(adata.uns.keys()),
            'layers':      list(adata.layers.keys()),
            'data_state':  _analyze_data_state(adata),
            # New high-performance metadata
            'schema': schema,
            'chunk_info': chunk_info,
            'summary': summary,
            'success': True
        }

        return jsonify(response_data)

    except Exception as e:
        logging.error(f"Upload failed: {e}")
        return jsonify({'error': str(e)}), 500


@bp.route('/upload_preview', methods=['POST'])
def upload_file_preview():
    """Upload h5ad file and open in backed='r' (preview / memory-efficient) mode."""
    if 'file' not in request.files:
        return jsonify({'error': 'No file provided'}), 400

    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No file selected'}), 400

    if not file.filename.endswith('.h5ad'):
        return jsonify({'error': 'File must be .h5ad format'}), 400

    try:
        import scanpy as sc
        filename = secure_filename(file.filename)
        filepath = os.path.join(bp.upload_folder, filename)
        file.save(filepath)

        # Open in backed read-only mode — does NOT load X into RAM
        adata_backed = sc.read_h5ad(filepath, backed='r')

        # Store in state so visualisation endpoints can use it
        bp.state.current_adata = adata_backed
        bp.state.current_filename = filename
        # Mark as preview (backed) so API tools can refuse to mutate it
        bp.state.is_preview_mode = True

        # Pass the already-opened backed adata to the adaptor (avoids re-loading fully)
        bp.state.current_adaptor = HighPerformanceAnndataAdaptor(adata_backed)

        embeddings = _canonical_embedding_keys(adata_backed)

        response_data = {
            'filename': filename,
            'n_cells':      adata_backed.n_obs,
            'n_genes':      adata_backed.n_vars,
            'embeddings':   embeddings,
            'obs_columns':  list(adata_backed.obs.columns),
            'var_columns':  list(adata_backed.var.columns),
            'uns_keys':     list(adata_backed.uns.keys()),
            'layers':       list(adata_backed.layers.keys()) if adata_backed.layers else [],
            'preview_mode': True,
            'success': True,
        }
        return jsonify(response_data)

    except Exception as e:
        logging.error(f"Preview upload failed: {e}")
        return jsonify({'error': str(e)}), 500


def _resolve_h5ad_path(raw_path, file_root):
    """Accept either an absolute path or a relative path (relative to file_root)."""
    from pathlib import Path as _Path
    p = _Path(raw_path)
    if p.is_absolute():
        return p.resolve()
    return resolve_browse_path(file_root, raw_path)


@bp.route('/load_from_server', methods=['POST'])
def load_from_server():
    """Load h5ad file from a server-side path (analysis/full mode)."""
    data = request.json if request.json else {}
    rel_path = data.get('path', '')
    if not rel_path:
        return jsonify({'error': 'No path provided'}), 400

    try:
        target = _resolve_h5ad_path(rel_path, bp.state.file_root)
    except ValueError:
        return jsonify({'error': 'Invalid path'}), 400

    if not target.exists() or not target.is_file():
        return jsonify({'error': 'File not found'}), 404

    if not is_h5ad_file(target):
        return jsonify({'error': 'File must be .h5ad format'}), 400

    try:
        filepath = str(target)
        filename = target.name

        bp.state.current_adaptor = HighPerformanceAnndataAdaptor(filepath)
        bp.state.current_adata = bp.state.current_adaptor.adata
        bp.state.current_filename = filename
        bp.state.is_preview_mode = False

        if 'X_random' not in bp.state.current_adata.obsm:
            seed = abs(hash(filepath)) % (2**32 - 1)
            rng = np.random.RandomState(seed)
            bp.state.current_adata.obsm['X_random'] = rng.rand(
                bp.state.current_adata.n_obs, 2).astype(np.float32)

        try:
            bp.state.kernel_executor.sync_adata(bp.state.current_adata)
        except Exception:
            pass

        from utils.adata_helpers import analyze_data_state as _analyze_data_state
        adata = bp.state.current_adata
        embeddings = _canonical_embedding_keys(adata)
        response_data = {
            'filename': filename,
            'n_cells': adata.n_obs,
            'n_genes': adata.n_vars,
            'embeddings': embeddings,
            'obs_columns': list(adata.obs.columns),
            'var_columns': list(adata.var.columns),
            'uns_keys': list(adata.uns.keys()),
            'layers': list(adata.layers.keys()),
            'data_state': _analyze_data_state(adata),
            'success': True
        }
        return jsonify(response_data)

    except Exception as e:
        logging.error(f"Load from server failed: {e}")
        return jsonify({'error': str(e)}), 500


@bp.route('/load_preview_from_server', methods=['POST'])
def load_preview_from_server():
    """Load h5ad file from a server-side path (preview/backed mode)."""
    data = request.json if request.json else {}
    rel_path = data.get('path', '')
    if not rel_path:
        return jsonify({'error': 'No path provided'}), 400

    try:
        target = _resolve_h5ad_path(rel_path, bp.state.file_root)
    except ValueError:
        return jsonify({'error': 'Invalid path'}), 400

    if not target.exists() or not target.is_file():
        return jsonify({'error': 'File not found'}), 404

    if not is_h5ad_file(target):
        return jsonify({'error': 'File must be .h5ad format'}), 400

    try:
        import scanpy as sc
        filepath = str(target)
        filename = target.name

        adata_backed = sc.read_h5ad(filepath, backed='r')
        bp.state.current_adata = adata_backed
        bp.state.current_filename = filename
        bp.state.is_preview_mode = True
        bp.state.current_adaptor = HighPerformanceAnndataAdaptor(adata_backed)

        embeddings = _canonical_embedding_keys(adata_backed)
        response_data = {
            'filename': filename,
            'n_cells': adata_backed.n_obs,
            'n_genes': adata_backed.n_vars,
            'embeddings': embeddings,
            'obs_columns': list(adata_backed.obs.columns),
            'var_columns': list(adata_backed.var.columns),
            'uns_keys': list(adata_backed.uns.keys()),
            'layers': list(adata_backed.layers.keys()) if adata_backed.layers else [],
            'preview_mode': True,
            'success': True,
        }
        return jsonify(response_data)

    except Exception as e:
        logging.error(f"Load preview from server failed: {e}")
        return jsonify({'error': str(e)}), 500


@bp.route('/schema', methods=['GET'])
def get_schema():
    """Get data schema."""
    if bp.state.current_adaptor is None:
        return jsonify({'error': 'No data loaded'}), 400

    try:
        schema = bp.state.current_adaptor.get_schema()
        return jsonify({'schema': schema})
    except Exception as e:
        logging.error(f"Schema retrieval failed: {e}")
        return jsonify({'error': str(e)}), 500


@bp.route('/data/obs', methods=['GET'])
def get_obs_data():
    """Get observation annotations as FlatBuffers."""
    if bp.state.current_adaptor is None:
        return jsonify({'error': 'No data loaded'}), 400

    try:
        columns = request.args.getlist('columns')
        chunk_index = request.args.get('chunk', 0, type=int)

        if chunk_index > 0:
            # Chunked data request
            fbs_data = bp.state.current_adaptor.get_chunked_data('obs', chunk_index, columns=columns if columns else None)
        else:
            # Full data request
            fbs_data = bp.state.current_adaptor.get_obs_fbs(columns if columns else None)

        response = make_response(fbs_data)
        response.headers['Content-Type'] = 'application/octet-stream'
        response.headers['Cache-Control'] = 'public, max-age=3600'
        return response

    except Exception as e:
        logging.error(f"Obs data retrieval failed: {e}")
        return jsonify({'error': str(e)}), 500


@bp.route('/data/embedding/<embedding_name>', methods=['GET'])
def get_embedding_data(embedding_name):
    """Get embedding coordinates as FlatBuffers."""
    if bp.state.current_adaptor is None:
        return jsonify({'error': 'No data loaded'}), 400

    try:
        chunk_index = request.args.get('chunk', 0, type=int)
        out_format = request.args.get('format', 'fbs')

        if chunk_index > 0:
            if out_format == 'json':
                # JSON slice response
                if embedding_name == 'random' or embedding_name == 'X_random':
                    coords = bp.state.current_adaptor._get_random_embedding()
                else:
                    embedding_key = _resolve_embedding_key(bp.state.current_adaptor.adata, embedding_name)
                    coords = bp.state.current_adaptor.adata.obsm[embedding_key]
                start = chunk_index * bp.state.current_adaptor.chunk_size
                end = min(start + bp.state.current_adaptor.chunk_size, bp.state.current_adaptor.n_obs)
                return jsonify({'x': coords[start:end, 0].tolist(), 'y': coords[start:end, 1].tolist()})
            else:
                # FBS slice
                fbs_data = bp.state.current_adaptor.get_chunked_data('embedding', chunk_index, embedding_name=embedding_name)
                response = make_response(fbs_data)
                response.headers['Content-Type'] = 'application/octet-stream'
                response.headers['Cache-Control'] = 'public, max-age=3600'
                return response
        else:
            if out_format == 'json':
                if embedding_name == 'random' or embedding_name == 'X_random':
                    coords = bp.state.current_adaptor._get_random_embedding()
                else:
                    embedding_key = _resolve_embedding_key(bp.state.current_adaptor.adata, embedding_name)
                    coords = bp.state.current_adaptor.adata.obsm[embedding_key]
                return jsonify({'x': coords[:, 0].tolist(), 'y': coords[:, 1].tolist()})
            else:
                fbs_data = bp.state.current_adaptor.get_embedding_fbs(embedding_name)
                response = make_response(fbs_data)
                response.headers['Content-Type'] = 'application/octet-stream'
                response.headers['Cache-Control'] = 'public, max-age=3600'
                return response

    except KeyError as e:
        return jsonify({'error': f'Embedding not found: {embedding_name}'}), 404
    except Exception as e:
        logging.error(f"Embedding data retrieval failed: {e}")
        return jsonify({'error': str(e)}), 500


@bp.route('/data/expression', methods=['POST'])
def get_expression_data():
    """Get gene expression data as FlatBuffers."""
    if bp.state.current_adaptor is None:
        return jsonify({'error': 'No data loaded'}), 400

    try:
        data = request.json
        gene_names = data.get('genes', [])
        cell_indices = data.get('cell_indices', None)

        if not gene_names:
            return jsonify({'error': 'No genes specified'}), 400

        fbs_data = bp.state.current_adaptor.get_expression_fbs(gene_names, cell_indices)

        response = make_response(fbs_data)
        response.headers['Content-Type'] = 'application/octet-stream'
        return response

    except ValueError as e:
        return jsonify({'error': str(e)}), 404
    except Exception as e:
        logging.error(f"Expression data retrieval failed: {e}")
        return jsonify({'error': str(e)}), 500


@bp.route('/save', methods=['POST'])
def save_data():
    """Save processed data as h5ad file."""
    if bp.state.current_adata is None:
        return jsonify({'error': 'No data to save'}), 400

    try:
        # Save to temporary file
        output_path = os.path.join(bp.upload_folder, f'processed_{bp.state.current_filename}')
        bp.state.current_adata.write_h5ad(output_path)

        return send_file(output_path, as_attachment=True,
                        download_name=f'processed_{bp.state.current_filename}')

    except Exception as e:
        return jsonify({'error': str(e)}), 500


@bp.route('/data_info', methods=['GET'])
def get_data_info():
    """Get detailed information about the loaded data."""
    if bp.state.current_adata is None:
        return jsonify({'error': 'No data loaded'}), 400

    try:
        info = {
            'shape': list(bp.state.current_adata.shape),
            'obs_keys': list(bp.state.current_adata.obs.keys()),
            'var_keys': list(bp.state.current_adata.var.keys()),
            'obsm_keys': list(bp.state.current_adata.obsm.keys()) if bp.state.current_adata.obsm else [],
            'uns_keys': list(bp.state.current_adata.uns.keys()) if bp.state.current_adata.uns else [],
            'layers': list(bp.state.current_adata.layers.keys()) if bp.state.current_adata.layers else []
        }

        # Add some statistics
        if hasattr(bp.state.current_adata.X, 'toarray'):
            X = bp.state.current_adata.X.toarray()
        else:
            X = bp.state.current_adata.X

        info['statistics'] = {
            'mean_genes_per_cell': float(np.mean(np.sum(X > 0, axis=1))),
            'mean_counts_per_cell': float(np.mean(np.sum(X, axis=1))),
            'total_counts': float(np.sum(X))
        }

        return jsonify(info)

    except Exception as e:
        logging.error(f"Data info retrieval failed: {e}")
        return jsonify({'error': str(e)}), 500


@bp.route('/data_status', methods=['GET'])
def get_data_status():
    """Get current data status (basic). Use /api/status for full UI-compatible response."""
    return jsonify({
        'loaded': bp.state.current_adata is not None,
        'filename': bp.state.current_filename,
        'cells': bp.state.current_adata.n_obs if bp.state.current_adata else 0,
        'genes': bp.state.current_adata.n_vars if bp.state.current_adata else 0,
    })


@bp.route('/genes', methods=['GET'])
def get_genes():
    """Get all gene names."""
    if bp.state.current_adata is None:
        return jsonify({'error': 'No data loaded'}), 400

    try:
        genes = list(bp.state.current_adata.var_names)
        return jsonify({'genes': genes})
    except Exception as e:
        return jsonify({'error': str(e)}), 500


@bp.route('/gene_search', methods=['POST'])
def search_genes():
    """Search for genes by name pattern."""
    if bp.state.current_adata is None:
        return jsonify({'error': 'No data loaded'}), 400

    try:
        data = request.json
        pattern = data.get('pattern', '').upper()

        if not pattern:
            return jsonify({'genes': []})

        # Search genes that match the pattern
        matching_genes = [gene for gene in bp.state.current_adata.var_names if pattern in gene.upper()]

        # Limit results to avoid too many matches
        matching_genes = matching_genes[:20]

        return jsonify({'genes': matching_genes})

    except Exception as e:
        return jsonify({'error': str(e)}), 500


@bp.route('/qc_prefixes', methods=['GET'])
def get_qc_prefixes():
    """Detect common QC gene prefixes (e.g., mitochondrial) from var_names."""
    if bp.state.current_adata is None:
        return jsonify({'error': 'No data loaded'}), 400

    try:
        import re
        var_names = bp.state.current_adata.var_names.astype(str)
        prefixes = set()

        # Detect typical mt prefixes
        for name in var_names[: min(10000, len(var_names))]:  # sample to be safe
            if len(name) < 3:
                continue
            if re.match(r'^(mt|MT|Mt|mT)+', name):
                prefixes.add(name[:3])  # e.g., 'MT-' or 'mt-'

        # If underscore variants like 'mt_' present
        for name in var_names[: min(60000, len(var_names))]:
            if name.lower().startswith('mt'):
                prefixes.add(name[:3])

        # Fallback to common defaults
        if not prefixes:
            prefixes.update(['MT-', 'mt-'])

        return jsonify({'mt_prefixes': sorted(prefixes)})
    except Exception as e:
        return jsonify({'error': str(e)}), 500


@bp.route('/filter', methods=['POST'])
def filter_cells():
    """Filter cells based on criteria."""
    if bp.state.current_adaptor is None:
        return jsonify({'error': 'No data loaded'}), 400

    try:
        filters = request.json
        indices = bp.state.current_adaptor.filter_cells(filters)

        return jsonify({
            'filtered_indices': indices,
            'count': len(indices)
        })

    except Exception as e:
        logging.error(f"Cell filtering failed: {e}")
        return jsonify({'error': str(e)}), 500


@bp.route('/differential_expression', methods=['POST'])
def compute_differential_expression():
    """Compute differential expression between groups."""
    if bp.state.current_adaptor is None:
        return jsonify({'error': 'No data loaded'}), 400

    try:
        data = request.json
        group1_indices = data.get('group1_indices', [])
        group2_indices = data.get('group2_indices', [])
        method = data.get('method', 'wilcoxon')
        n_genes = data.get('n_genes', 100)

        if not group1_indices or not group2_indices:
            return jsonify({'error': 'Both groups must have cells'}), 400

        # Run in thread pool to avoid blocking
        future = bp.state.thread_pool.submit(
            bp.state.current_adaptor.get_differential_expression,
            group1_indices, group2_indices, method, n_genes
        )

        result = future.result(timeout=30)  # 30 second timeout

        return jsonify(result)

    except Exception as e:
        logging.error(f"Differential expression failed: {e}")
        return jsonify({'error': str(e)}), 500


@bp.route('/export_plot_data', methods=['POST'])
def export_plot_data():
    """Export current plot data as CSV."""
    if bp.state.current_adata is None:
        return jsonify({'error': 'No data loaded'}), 400

    try:
        data = request.json
        embedding = data.get('embedding', '')
        color_by = data.get('color_by', '')

        # Create CSV with embedding coordinates and color data
        import pandas as pd
        import io

        # Get embedding coordinates
        embedding_key = f'X_{embedding}' if not embedding.startswith('X_') else embedding
        if embedding_key not in bp.state.current_adata.obsm:
            return jsonify({'error': f'Embedding {embedding} not found'}), 404

        coords = bp.state.current_adata.obsm[embedding_key]
        df = pd.DataFrame({
            f'{embedding}_1': coords[:, 0],
            f'{embedding}_2': coords[:, 1]
        })

        # Add color data if specified
        if color_by and color_by in bp.state.current_adata.obs.columns:
            df[color_by] = bp.state.current_adata.obs[color_by].values

        # Convert to CSV
        csv_buffer = io.StringIO()
        df.to_csv(csv_buffer, index=True)
        csv_buffer.seek(0)

        from flask import Response
        return Response(
            csv_buffer.getvalue(),
            mimetype='text/csv',
            headers={'Content-Disposition': f'attachment; filename=plot_data_{embedding}.csv'}
        )

    except Exception as e:
        logging.error(f"Export plot data failed: {e}")
        return jsonify({'error': str(e)}), 500


# Initialize blueprint with dependencies (will be set by app.py)
bp.state = None
bp.upload_folder = None
