from flask import Flask, jsonify, send_from_directory, request, send_file
from flask_cors import CORS
import os
import re
import h5py
import numpy as np
from datetime import datetime, timedelta
import base64
import struct
import traceback
import json
import io
import mimetypes
from pyproj import Transformer
import shutil

app = Flask(__name__, static_folder='static', static_url_path='')
CORS(app)

DATA_DIR = os.environ.get('DATA_DIR', '/data')
LOCATIONS_FILE = '/home/floppyrj45/seismic-viewer/locations.json'

# ========== UTILITY FUNCTIONS ==========

def load_locations():
    """Load locations from JSON file."""
    if os.path.exists(LOCATIONS_FILE):
        try:
            with open(LOCATIONS_FILE, 'r') as f:
                return json.load(f)
        except:
            return {}
    return {}

def save_locations(locations):
    """Save locations to JSON file."""
    with open(LOCATIONS_FILE, 'w') as f:
        json.dump(locations, f, indent=2)

def get_h5_metadata(filepath):
    """Extract comprehensive metadata from H5 file."""
    metadata = {}
    try:
        with h5py.File(filepath, 'r') as f:
            # File-level attributes
            for k, v in f.attrs.items():
                try:
                    metadata[k] = v.item() if hasattr(v, 'item') else str(v)
                except:
                    metadata[k] = str(v)
            
            # Metadata group attributes
            if 'metadata' in f:
                for k, v in f['metadata'].attrs.items():
                    try:
                        metadata[k] = v.item() if hasattr(v, 'item') else str(v)
                    except:
                        metadata[k] = str(v)
            
            # Calibration group attributes
            if 'calibration' in f:
                calibration = {}
                for k, v in f['calibration'].attrs.items():
                    try:
                        calibration[k] = v.item() if hasattr(v, 'item') else str(v)
                    except:
                        calibration[k] = str(v)
                metadata['calibration'] = calibration
            
            # Channel information
            channels = []
            for group_name in ['calibrated_data', 'raw_data']:
                if group_name in f:
                    for ds_name in f[group_name].keys():
                        ds = f[group_name][ds_name]
                        ch_info = {
                            'name': f"{group_name}/{ds_name}",
                            'shape': list(ds.shape),
                            'dtype': str(ds.dtype),
                            'samples': int(ds.shape[0]) if len(ds.shape) > 0 else 0
                        }
                        for k, v in ds.attrs.items():
                            try:
                                ch_info[k] = v.item() if hasattr(v, 'item') else str(v)
                            except:
                                ch_info[k] = str(v)
                        channels.append(ch_info)
            metadata['channels'] = channels
            
            # Compute timestamps
            if 'creation_date' in metadata:
                try:
                    start_time = datetime.fromisoformat(metadata['creation_date'].replace('Z', '+00:00'))
                    metadata['start_timestamp'] = start_time.isoformat()
                    
                    if 'duration_sec' in metadata:
                        duration_sec = float(metadata['duration_sec'])
                        end_time = start_time + timedelta(seconds=duration_sec)
                        metadata['end_timestamp'] = end_time.isoformat()
                        
                        # Human-readable duration
                        hours = int(duration_sec // 3600)
                        minutes = int((duration_sec % 3600) // 60)
                        seconds = int(duration_sec % 60)
                        metadata['duration_human'] = f"{hours}h {minutes}m {seconds}s"
                except Exception as e:
                    metadata['timestamp_error'] = str(e)
            
            return metadata
    except Exception as e:
        return {'error': str(e)}

# ========== ROUTES ==========

@app.route('/')
def index():
    return send_from_directory('static', 'index.html')

@app.route('/api/files')
def list_files():
    try:
        files = []
        for fn in os.listdir(DATA_DIR):
            if fn.endswith(('.segy', '.sgy', '.h5')):
                fp = os.path.join(DATA_DIR, fn)
                stat = os.stat(fp)
                ftype = 'h5' if fn.endswith('.h5') else 'segy'
                files.append({
                    'name': fn, 'size': stat.st_size,
                    'modified': datetime.fromtimestamp(stat.st_mtime).isoformat(),
                    'type': ftype
                })
        files.sort(key=lambda x: x['name'])
        return jsonify({'files': files, 'count': len(files)})
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/file/<filename>/info')
def file_info(filename):
    try:
        fp = os.path.join(DATA_DIR, filename)
        if not os.path.exists(fp):
            return jsonify({'error': 'Not found'}), 404
        stat = os.stat(fp)
        info = {
            'name': filename,
            'size': stat.st_size,
            'size_mb': round(stat.st_size / 1048576, 2),
            'modified': datetime.fromtimestamp(stat.st_mtime).isoformat(),
            'type': 'h5' if filename.endswith('.h5') else 'segy',
        }
        if filename.endswith('.h5'):
            with h5py.File(fp, 'r') as f:
                datasets = []
                def visitor(name, obj):
                    if isinstance(obj, h5py.Dataset):
                        datasets.append({
                            'path': name,
                            'shape': list(obj.shape),
                            'dtype': str(obj.dtype),
                            'samples': int(obj.shape[0]) if len(obj.shape) > 0 else 0
                        })
                f.visititems(visitor)
                info['datasets'] = datasets
                info['num_datasets'] = len(datasets)
                cal = [d for d in datasets if 'calibrated' in d['path']]
                raw = [d for d in datasets if 'raw' in d['path']]
                info['calibrated_channels'] = len(cal)
                info['raw_channels'] = len(raw)
                if cal:
                    info['samples_per_channel'] = cal[0]['samples']
                
                # Add sample rate and duration for sidebar display
                if 'metadata' in f:
                    sample_rate = f['metadata'].attrs.get('sample_rate_hz', None)
                    if sample_rate and cal:
                        total_samples = cal[0]['samples']
                        duration_sec = total_samples / sample_rate
                        hours = int(duration_sec // 3600)
                        minutes = int((duration_sec % 3600) // 60)
                        info['duration_human'] = f"{hours}h{minutes:02d}m"
                        info['sample_rate_hz'] = sample_rate
                        info['num_channels'] = len(cal)
        elif filename.endswith(('.segy', '.sgy')):
            try:
                from obspy import read as obspy_read
                st = obspy_read(fp, headonly=True)
                info['num_traces'] = len(st)
                if len(st) > 0:
                    info['samples_per_trace'] = st[0].stats.npts
                    info['sample_rate'] = st[0].stats.sampling_rate
                    info['delta'] = st[0].stats.delta
            except Exception as e:
                info['segy_error'] = str(e)
        return jsonify(info)
    except Exception as e:
        return jsonify({'error': str(e)}), 500

# ========== NEW: METADATA ENDPOINT ==========

@app.route('/api/file/<filename>/metadata')
def file_metadata(filename):
    """Return detailed metadata for a file."""
    try:
        fp = os.path.join(DATA_DIR, filename)
        if not os.path.exists(fp):
            return jsonify({'error': 'Not found'}), 404
        
        stat = os.stat(fp)
        result = {
            'filename': filename,
            'file_size_bytes': stat.st_size,
            'file_size_mb': round(stat.st_size / 1048576, 2),
            'modified': datetime.fromtimestamp(stat.st_mtime).isoformat(),
        }
        
        if filename.endswith('.h5'):
            metadata = get_h5_metadata(fp)
            result.update(metadata)
        elif filename.endswith(('.segy', '.sgy')):
            try:
                from obspy import read as obspy_read
                st = obspy_read(fp, headonly=True)
                result['num_traces'] = len(st)
                if len(st) > 0:
                    result['sample_rate'] = st[0].stats.sampling_rate
                    result['samples_per_trace'] = st[0].stats.npts
                    result['start_timestamp'] = st[0].stats.starttime.isoformat()
                    result['end_timestamp'] = st[0].stats.endtime.isoformat()
                    duration_sec = (st[0].stats.endtime - st[0].stats.starttime)
                    hours = int(duration_sec // 3600)
                    minutes = int((duration_sec % 3600) // 60)
                    seconds = int(duration_sec % 60)
                    result['duration_human'] = f"{hours}h {minutes}m {seconds}s"
                    result['duration_sec'] = duration_sec
            except Exception as e:
                result['segy_error'] = str(e)
        
        return jsonify(result)
    except Exception as e:
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500

# ========== NEW: LOCATIONS ENDPOINT ==========

@app.route('/api/locations', methods=['GET'])
def get_locations():
    """Get all file locations."""
    try:
        locations = load_locations()
        return jsonify(locations)
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/locations', methods=['POST'])
def update_location():
    """Update location for a file."""
    try:
        data = request.json
        filename = data.get('filename')
        lat = data.get('lat')
        lon = data.get('lon')
        
        if not all([filename, lat is not None, lon is not None]):
            return jsonify({'error': 'Missing filename, lat, or lon'}), 400
        
        locations = load_locations()
        locations[filename] = {
            'lat': float(lat),
            'lon': float(lon),
            'updated': datetime.now().isoformat()
        }
        save_locations(locations)
        
        return jsonify({'success': True, 'locations': locations})
    except Exception as e:
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500

# ========== NEW: EXTRACT ENDPOINT ==========

@app.route('/api/file/<filename>/extract')
def extract_segment(filename):
    """Extract a time segment from a file and return as CSV or H5."""
    try:
        fp = os.path.join(DATA_DIR, filename)
        if not os.path.exists(fp):
            return jsonify({'error': 'Not found'}), 404
        
        # Parameters
        start_sec = float(request.args.get('start_sec', 0))
        duration_sec = float(request.args.get('duration_sec', 1200))  # Default 20 min
        channel = request.args.get('channel', 'calibrated_data/channel_1')
        format_type = request.args.get('format', 'csv')  # csv or h5
        
        if not filename.endswith('.h5'):
            return jsonify({'error': 'Extract only supported for H5 files'}), 400
        
        with h5py.File(fp, 'r') as f:
            if channel not in f:
                return jsonify({'error': f'Channel {channel} not found'}), 404
            
            # Get metadata
            sample_rate = f['metadata'].attrs.get('sample_rate_hz', 500)
            total_samples = f[channel].shape[0]
            
            # Calculate sample range
            start_sample = int(start_sec * sample_rate)
            end_sample = int((start_sec + duration_sec) * sample_rate)
            
            # Clamp to valid range
            start_sample = max(0, min(start_sample, total_samples - 1))
            end_sample = max(start_sample + 1, min(end_sample, total_samples))
            
            # Extract data
            data = f[channel][start_sample:end_sample]
            
            if format_type == 'csv':
                # Generate CSV
                output = io.StringIO()
                output.write(f"# Channel: {channel}\n")
                output.write(f"# Start: {start_sec}s, Duration: {duration_sec}s\n")
                output.write(f"# Sample rate: {sample_rate} Hz\n")
                output.write(f"# Samples: {len(data)}\n")
                output.write("sample,time_sec,value\n")
                
                for i, val in enumerate(data):
                    time_s = (start_sample + i) / sample_rate
                    output.write(f"{start_sample + i},{time_s:.6f},{val}\n")
                
                output.seek(0)
                return send_file(
                    io.BytesIO(output.getvalue().encode('utf-8')),
                    mimetype='text/csv',
                    as_attachment=True,
                    download_name=f'{filename.replace(".h5", "")}_extract_{int(start_sec)}_{int(duration_sec)}s.csv'
                )
            
            elif format_type == 'h5':
                # Generate H5
                output = io.BytesIO()
                with h5py.File(output, 'w') as out_f:
                    # Copy metadata
                    if 'metadata' in f:
                        meta_group = out_f.create_group('metadata')
                        for k, v in f['metadata'].attrs.items():
                            meta_group.attrs[k] = v
                        meta_group.attrs['extract_start_sec'] = start_sec
                        meta_group.attrs['extract_duration_sec'] = duration_sec
                        meta_group.attrs['extract_start_sample'] = start_sample
                        meta_group.attrs['extract_end_sample'] = end_sample
                    
                    # Create dataset
                    ds = out_f.create_dataset('data', data=data, compression='gzip')
                    ds.attrs['channel'] = channel
                    ds.attrs['sample_rate_hz'] = sample_rate
                    
                    # Copy channel attributes
                    if channel in f:
                        for k, v in f[channel].attrs.items():
                            ds.attrs[k] = v
                
                output.seek(0)
                return send_file(
                    output,
                    mimetype='application/x-hdf5',
                    as_attachment=True,
                    download_name=f'{filename.replace(".h5", "")}_extract_{int(start_sec)}_{int(duration_sec)}s.h5'
                )
            
            else:
                return jsonify({'error': 'Invalid format. Use csv or h5'}), 400
    
    except Exception as e:
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500

# ========== EXISTING ENDPOINTS (kept for compatibility) ==========

@app.route('/api/file/<filename>/section')
def file_section(filename):
    """Return 2D section data as base64 float32 array."""
    try:
        fp = os.path.join(DATA_DIR, filename)
        if not os.path.exists(fp):
            return jsonify({'error': 'Not found'}), 404

        trace_start = int(request.args.get('trace_start', 0))
        trace_count = int(request.args.get('trace_count', 200))
        samples_per_trace = int(request.args.get('samples_per_trace', 1000))
        channel = request.args.get('channel', 'calibrated_data/channel_1')
        max_samples = int(request.args.get('max_samples', 500))

        if filename.endswith('.h5'):
            with h5py.File(fp, 'r') as f:
                if channel not in f:
                    avail = []
                    f.visititems(lambda n, o: avail.append(n) if isinstance(o, h5py.Dataset) else None)
                    return jsonify({'error': f'Channel not found', 'available': avail}), 404
                
                data = f[channel]
                total_samples = data.shape[0]
                total_traces = total_samples // samples_per_trace
                
                trace_start = max(0, min(trace_start, max(0, total_traces - 1)))
                trace_end = min(trace_start + trace_count, total_traces)
                actual_count = trace_end - trace_start
                
                if actual_count <= 0:
                    return jsonify({'error': 'No traces in range'}), 400
                
                start_sample = trace_start * samples_per_trace
                end_sample = trace_end * samples_per_trace
                raw = data[start_sample:end_sample]
                
                section = raw.reshape(actual_count, samples_per_trace).astype(np.float32)
                
                if samples_per_trace > max_samples:
                    factor = samples_per_trace // max_samples
                    trimmed = section[:, :factor * max_samples]
                    section = trimmed.reshape(actual_count, max_samples, factor).mean(axis=2).astype(np.float32)
                    actual_spt = max_samples
                else:
                    actual_spt = samples_per_trace

                vmin = float(np.min(section))
                vmax = float(np.max(section))
                encoded = base64.b64encode(section.tobytes()).decode('ascii')
                
                return jsonify({
                    'num_traces': int(actual_count),
                    'samples_per_trace': actual_spt,
                    'trace_start': trace_start,
                    'total_traces': total_traces,
                    'total_samples': int(total_samples),
                    'vmin': vmin, 'vmax': vmax,
                    'data_b64': encoded,
                    'dtype': 'float32'
                })

        elif filename.endswith(('.segy', '.sgy')):
            try:
                from obspy import read as obspy_read
                st = obspy_read(fp)
                total_traces = len(st)
                
                trace_start = max(0, min(trace_start, total_traces - 1))
                trace_end = min(trace_start + trace_count, total_traces)
                actual_count = trace_end - trace_start
                
                if actual_count <= 0:
                    return jsonify({'error': 'No traces'}), 400
                
                spt = st[0].stats.npts
                actual_spt = spt
                
                section = np.zeros((actual_count, spt), dtype=np.float32)
                for i in range(actual_count):
                    tr = st[trace_start + i]
                    n = min(spt, tr.stats.npts)
                    section[i, :n] = tr.data[:n].astype(np.float32)
                
                if spt > max_samples:
                    factor = spt // max_samples
                    trimmed = section[:, :factor * max_samples]
                    section = trimmed.reshape(actual_count, max_samples, factor).mean(axis=2).astype(np.float32)
                    actual_spt = max_samples
                
                vmin = float(np.min(section))
                vmax = float(np.max(section))
                encoded = base64.b64encode(section.tobytes()).decode('ascii')
                
                return jsonify({
                    'num_traces': int(actual_count),
                    'samples_per_trace': actual_spt,
                    'trace_start': trace_start,
                    'total_traces': total_traces,
                    'total_samples': int(total_traces * spt),
                    'vmin': vmin, 'vmax': vmax,
                    'data_b64': encoded,
                    'dtype': 'float32'
                })
            except Exception as e:
                return jsonify({'error': f'SEGY read error: {e}'}), 500
        
        return jsonify({'error': 'Unsupported format'}), 400
    except Exception as e:
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500

@app.route('/api/file/<filename>/headers')
def file_headers(filename):
    """Return trace headers for SEGY or dataset attributes for H5."""
    try:
        fp = os.path.join(DATA_DIR, filename)
        if not os.path.exists(fp):
            return jsonify({'error': 'Not found'}), 404
        
        if filename.endswith('.h5'):
            with h5py.File(fp, 'r') as f:
                attrs = {}
                for k, v in f.attrs.items():
                    try:
                        attrs[k] = v.item() if hasattr(v, 'item') else str(v)
                    except:
                        attrs[k] = str(v)
                
                datasets = []
                def visitor(name, obj):
                    if isinstance(obj, h5py.Dataset):
                        ds_info = {
                            'path': name,
                            'shape': list(obj.shape),
                            'dtype': str(obj.dtype),
                            'size': int(np.prod(obj.shape)),
                        }
                        for ak, av in obj.attrs.items():
                            try:
                                ds_info[f'attr:{ak}'] = av.item() if hasattr(av, 'item') else str(av)
                            except:
                                ds_info[f'attr:{ak}'] = str(av)
                        datasets.append(ds_info)
                f.visititems(visitor)
                
                groups = []
                def gvisitor(name, obj):
                    if isinstance(obj, h5py.Group):
                        g_info = {'path': name}
                        for ak, av in obj.attrs.items():
                            try:
                                g_info[f'attr:{ak}'] = av.item() if hasattr(av, 'item') else str(av)
                            except:
                                g_info[f'attr:{ak}'] = str(av)
                        if len(g_info) > 1:
                            groups.append(g_info)
                f.visititems(gvisitor)
                
                return jsonify({
                    'type': 'h5',
                    'file_attrs': attrs,
                    'datasets': datasets,
                    'groups': groups
                })
        
        elif filename.endswith(('.segy', '.sgy')):
            try:
                from obspy import read as obspy_read
                from obspy.io.segy.segy import _read_segy
                
                segy = _read_segy(fp)
                
                bfh = {}
                if hasattr(segy, 'binary_file_header'):
                    h = segy.binary_file_header
                    for attr in ['job_identification_number', 'line_number', 'reel_number',
                                 'number_of_data_traces_per_ensemble', 'number_of_auxiliary_traces_per_ensemble',
                                 'sample_interval_in_microseconds', 'number_of_samples_per_data_trace',
                                 'data_sample_format_code']:
                        if hasattr(h, attr):
                            bfh[attr] = getattr(h, attr)
                
                trace_headers = []
                max_traces = min(50, len(segy.traces))
                for i in range(max_traces):
                    th = segy.traces[i].header
                    trace_headers.append({
                        'trace': i,
                        'inline': getattr(th, 'trace_sequence_number_within_line', None),
                        'crossline': getattr(th, 'trace_sequence_number_within_segy_file', None),
                        'field_record': getattr(th, 'original_field_record_number', None),
                        'trace_number': getattr(th, 'trace_number_within_the_original_field_record', None),
                        'cdp': getattr(th, 'ensemble_number', None),
                        'trace_in_ensemble': getattr(th, 'trace_number_within_the_ensemble', None),
                        'offset': getattr(th, 'distance_from_center_of_the_source_point_to_the_center_of_the_receiver_group', None),
                        'source_x': getattr(th, 'source_coordinate_x', None),
                        'source_y': getattr(th, 'source_coordinate_y', None),
                        'group_x': getattr(th, 'group_coordinate_x', None),
                        'group_y': getattr(th, 'group_coordinate_y', None),
                        'num_samples': getattr(th, 'number_of_samples_in_this_trace', None),
                        'sample_interval': getattr(th, 'sample_interval_in_ms_for_this_trace', None),
                    })
                
                return jsonify({
                    'type': 'segy',
                    'binary_header': bfh,
                    'trace_headers': trace_headers,
                    'total_traces': len(segy.traces)
                })
            except Exception as e:
                return jsonify({'error': f'SEGY header read error: {e}'}), 500
        
        return jsonify({'error': 'Unsupported'}), 400
    except Exception as e:
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500

@app.route('/api/file/<filename>/waveform')
def file_waveform(filename):
    try:
        fp = os.path.join(DATA_DIR, filename)
        if not os.path.exists(fp):
            return jsonify({'error': 'Not found'}), 404
        max_pts = int(request.args.get('points', 5000))
        channel = request.args.get('channel', 'calibrated_data/channel_1')
        if filename.endswith('.h5'):
            with h5py.File(fp, 'r') as f:
                if channel not in f:
                    return jsonify({'error': 'Channel not found'}), 404
                data = f[channel][:]
                n = len(data)
                if n <= max_pts:
                    values = data.tolist()
                    indices = list(range(n))
                else:
                    step = n // max_pts
                    indices = list(range(0, n, step))[:max_pts]
                    values = [float(data[i]) for i in indices]
                return jsonify({
                    'name': filename, 'channel': channel,
                    'total_samples': n,
                    'displayed_points': len(values),
                    'min': float(np.min(data)), 'max': float(np.max(data)),
                    'mean': float(np.mean(data)), 'std': float(np.std(data)),
                    'x': indices, 'y': values
                })
        return jsonify({'error': 'Waveform only for H5'}), 400
    except Exception as e:
        return jsonify({'error': str(e)}), 500


# ========== GEOSUP ENDPOINTS ==========

# UTM Zone 31N to WGS84 transformer
utm_to_wgs84 = Transformer.from_crs("EPSG:32631", "EPSG:4326", always_xy=True)

GEOSUP_DIR = os.path.join(app.static_folder, '')  # Static folder

def parse_sps_line(line):
    """Parse SPS S01/R01 format line. Fixed-width format."""
    if len(line) < 65:
        return None
    try:
        record_type = line[0].strip()  # S or R
        line_num = line[1:11].strip()
        point = line[11:21].strip()
        index = line[21:23].strip()
        # Easting and Northing are in columns ~26-52
        coord_part = line[46:66].strip()
        # Parse XXXXXX.XNNNNNNN.NN format
        match = re.match(r'(\d{5,6}\.\d{1,2})(\d{7}\.\d{1,2})', coord_part)
        if match:
            easting = float(match.group(1))
            northing = float(match.group(2))
        else:
            parts = coord_part.split('.')
            if len(parts) >= 2:
                easting_int = parts[0]
                rest = '.'.join(parts[1:])
                if len(rest) > 9:
                    easting = float(easting_int + '.' + rest[:2])
                    northing = float(rest[2:])
                else:
                    return None
            else:
                return None
        
        depth = line[66:72].strip() if len(line) > 52 else '0'
        depth = float(depth) if depth else 0
        
        return {
            'type': record_type,
            'line': line_num,
            'point': point,
            'index': index,
            'easting': easting,
            'northing': northing,
            'depth': depth
        }
    except Exception:
        return None

def parse_deployment_csv(filepath):
    """Parse deployment CSV file (SETE-0965P1002.csv format)."""
    records = []
    try:
        with open(filepath, 'r') as f:
            for line in f:
                parts = line.strip().split(',')
                if len(parts) >= 8:
                    try:
                        record = {
                            'line': parts[0],
                            'point': parts[1],
                            'index': parts[2],
                            'name': parts[3],
                            'easting': float(parts[5]),
                            'northing': float(parts[6]),
                            'depth': float(parts[7]),
                        }
                        if len(parts) >= 16:
                            record['year'] = int(parts[10])
                            record['day'] = int(parts[11])
                            record['hour'] = int(parts[12])
                            record['minute'] = int(parts[13])
                            record['second'] = int(parts[14])
                        records.append(record)
                    except (ValueError, IndexError):
                        continue
    except Exception:
        pass
    return records

@app.route('/api/geosup/positions')
def geosup_positions():
    """Return GeoJSON with source, receiver, and deployment positions."""
    try:
        features = []
        
        # Parse sources (S01)
        s01_path = os.path.join(GEOSUP_DIR, 'SeteSxPreplots.s01')
        if os.path.exists(s01_path):
            with open(s01_path, 'r') as f:
                for line in f:
                    record = parse_sps_line(line)
                    if record and record['type'] == 'S':
                        lon, lat = utm_to_wgs84.transform(record['easting'], record['northing'])
                        features.append({
                            'type': 'Feature',
                            'geometry': {'type': 'Point', 'coordinates': [lon, lat]},
                            'properties': {
                                'category': 'source',
                                'line': record['line'],
                                'point': record['point'],
                                'easting': record['easting'],
                                'northing': record['northing'],
                                'depth': record['depth']
                            }
                        })
        
        # Parse receivers (R01)
        r01_path = os.path.join(GEOSUP_DIR, 'SeteRxPreplots.r01')
        if os.path.exists(r01_path):
            with open(r01_path, 'r') as f:
                for line in f:
                    record = parse_sps_line(line)
                    if record and record['type'] == 'R':
                        lon, lat = utm_to_wgs84.transform(record['easting'], record['northing'])
                        features.append({
                            'type': 'Feature',
                            'geometry': {'type': 'Point', 'coordinates': [lon, lat]},
                            'properties': {
                                'category': 'receiver',
                                'line': record['line'],
                                'point': record['point'],
                                'easting': record['easting'],
                                'northing': record['northing'],
                                'depth': record['depth']
                            }
                        })
        
        # Parse deployments (CSV)
        csv_path = os.path.join(GEOSUP_DIR, 'SETE-0965P1002.csv')
        if os.path.exists(csv_path):
            deployments = parse_deployment_csv(csv_path)
            for d in deployments:
                lon, lat = utm_to_wgs84.transform(d['easting'], d['northing'])
                props = {
                    'category': 'deployment',
                    'line': d['line'],
                    'point': d['point'],
                    'name': d['name'],
                    'easting': d['easting'],
                    'northing': d['northing'],
                    'depth': d['depth']
                }
                if 'year' in d:
                    props['timestamp'] = f"{d['year']}-{d['day']:03d} {d['hour']:02d}:{d['minute']:02d}:{d['second']:02d}"
                features.append({
                    'type': 'Feature',
                    'geometry': {'type': 'Point', 'coordinates': [lon, lat]},
                    'properties': props
                })
        
        sources = len([f for f in features if f['properties']['category'] == 'source'])
        receivers = len([f for f in features if f['properties']['category'] == 'receiver'])
        deployments = len([f for f in features if f['properties']['category'] == 'deployment'])
        
        return jsonify({
            'type': 'FeatureCollection',
            'features': features,
            'counts': {'sources': sources, 'receivers': receivers, 'deployments': deployments, 'total': len(features)}
        })
    except Exception as e:
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500

@app.route('/api/geosup/documents')
def geosup_documents():
    """List all geosup documents in static folder."""
    try:
        documents = []
        doc_extensions = ('.pdf', '.docx', '.xlsx', '.doc', '.xls')
        
        for fn in os.listdir(GEOSUP_DIR):
            if fn.lower().endswith(doc_extensions):
                fp = os.path.join(GEOSUP_DIR, fn)
                stat = os.stat(fp)
                ext = os.path.splitext(fn)[1].lower()
                doc_type = {'.pdf': 'PDF', '.docx': 'Word', '.doc': 'Word', '.xlsx': 'Excel', '.xls': 'Excel'}.get(ext, 'Document')
                icon = {'.pdf': '📕', '.docx': '📘', '.doc': '📘', '.xlsx': '📗', '.xls': '📗'}.get(ext, '📄')
                is_gundalf = 'GUNDALF' in fn.upper()
                is_params = 'Acquisition_Parameters' in fn or 'SpiceRack' in fn
                
                documents.append({
                    'name': fn,
                    'size': stat.st_size,
                    'size_human': f"{stat.st_size / 1024:.1f} KB" if stat.st_size < 1048576 else f"{stat.st_size / 1048576:.1f} MB",
                    'type': doc_type,
                    'icon': icon,
                    'is_gundalf': is_gundalf,
                    'is_params': is_params,
                    'modified': datetime.fromtimestamp(stat.st_mtime).isoformat()
                })
        
        documents.sort(key=lambda d: (not d['is_gundalf'], not d['is_params'], d['name']))
        return jsonify({'documents': documents, 'count': len(documents)})
    except Exception as e:
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500

@app.route('/api/geosup/documents/<path:filename>')
def geosup_download(filename):
    """Download a geosup document."""
    try:
        safe_name = os.path.basename(filename)
        filepath = os.path.join(GEOSUP_DIR, safe_name)
        if not os.path.exists(filepath):
            return jsonify({'error': 'File not found'}), 404
        mime_type, _ = mimetypes.guess_type(filepath)
        return send_file(filepath, mimetype=mime_type or 'application/octet-stream', as_attachment=True, download_name=safe_name)
    except Exception as e:
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500

@app.route('/api/geosup/acquisition-params')
def geosup_acquisition_params():
    """Extract key parameters from the acquisition spreadsheet."""
    try:
        xlsx_path = os.path.join(GEOSUP_DIR, 'SBGS_Sete_SpiceRack_Acquisition_Parameters_Spreadsheet.xlsx')
        if not os.path.exists(xlsx_path):
            return jsonify({'error': 'Acquisition parameters file not found'}), 404
        from openpyxl import load_workbook
        wb = load_workbook(xlsx_path, data_only=True)
        params = {'sheets': [], 'summary': {}}
        for sheet_name in wb.sheetnames:
            ws = wb[sheet_name]
            sheet_data = {'name': sheet_name, 'rows': []}
            row_count = 0
            for row in ws.iter_rows(max_row=30, values_only=True):
                if any(cell is not None for cell in row):
                    sheet_data['rows'].append([str(cell) if cell is not None else '' for cell in row[:10]])
                    row_count += 1
                    if row_count >= 20:
                        break
            params['sheets'].append(sheet_data)
        wb.close()
        return jsonify(params)
    except Exception as e:
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500

@app.route('/api/pipeline-status')
def pipeline_status():
    """Return pipeline statistics and disk usage."""
    try:
        raw_dir = '/mnt/usb/Recordings'
        output_dir = '/mnt/usb/H5-Output'
        
        # Count files
        raw_count = 0
        if os.path.exists(raw_dir):
            raw_count = len([n for n in os.listdir(raw_dir) if n.lower().endswith(('.raw', '.manta'))])
            
        segy_count = 0
        h5_count = 0
        if os.path.exists(output_dir):
            files = os.listdir(output_dir)
            segy_count = len([n for n in files if n.lower().endswith(('.segy', '.sgy'))])
            h5_count = len([n for n in files if n.lower().endswith('.h5')])
            
        # Disk usage
        total, used, free = shutil.disk_usage(output_dir)
        
        # Conversion progress
        progress = None
        if os.path.exists('/tmp/segy2h5_progress.json'):
            try:
                with open('/tmp/segy2h5_progress.json', 'r') as f:
                    progress = json.load(f)
            except:
                pass
                
        return jsonify({
            'counts': {
                'raw': raw_count,
                'segy': segy_count,
                'h5': h5_count,
                'total_expected': 345  # Hardcoded target from context
            },
            'disk': {
                'total_gb': round(total / (1024**3), 2),
                'used_gb': round(used / (1024**3), 2),
                'free_gb': round(free / (1024**3), 2),
                'percent': round((used / total) * 100, 1)
            },
            'progress': progress
        })
    except Exception as e:
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=3001, debug=False)

@app.route('/test-map')
def test_map():
    return send_from_directory('static', 'test-map.html')