from pathlib import Path
import numpy as np

from .recording import Recording
from .Event import Event
from .RecordingMetadata import RecordingMetadata
from .breesy_scipy import _loadmat
from .errors import BreesyError, BreesyInternalError
from .load import _leave_only_eeg_channels
from .type_hints import enforce_type_hints


@enforce_type_hints
def load_eeglab(filename: str) -> Recording:
    """Load an EEGlab .set file (with optional .fdt companion file).

    :param filename: Path to the .set or .fdt file

    :return: Recording object containing the loaded data
    """
    if filename.endswith('.fdt'):
        filename = Path(filename).with_suffix('.set')
        if not filename.is_file():
            raise BreesyError('".fdt" filename provided, but it cannot be read without a corresponding ".set" file, which was not found.',
                              'Ensure that the ".set" file is in the same directory as your ".fdt" file. Their names should be the same, except for the extension.')

    set_data = _loadmat(filename)

    eeg = _extract_eeg_structure(set_data, filename)
    _validate_continuous_data(eeg, filename)

    # sample rate, channels, dimensions
    sample_rate, nbchan, pnts = _extract_metadata(eeg, filename)
    channel_names = _extract_channel_names(eeg, nbchan, filename)

    data = _load_eeglab_data(eeg, filename, nbchan, pnts)
    data, channel_names = _leave_only_eeg_channels(data, channel_names)
    events = _extract_events(eeg, sample_rate)

    metadata = RecordingMetadata(
        name=Path(filename).stem,
        file_path=Path(filename),
        description=f"EEGlab data loaded from {filename}"
    )

    return Recording(
        data=data,
        channel_names=channel_names,
        sample_rate=sample_rate,
        events=events,
        metadata=metadata
    )


def _extract_eeg_structure(set_data: dict, filename: str) -> dict:
    # Most EEGlab files have a top-level 'EEG' key
    if 'EEG' in set_data:
        eeg_raw = set_data['EEG']
        # Convert structured array to dict for easier access
        return _struct_to_dict(eeg_raw)

    # Fallback: look for meaningful keys (not MATLAB metadata)
    meaningful_keys = [k for k in set_data.keys() if not k.startswith('__')]

    if len(meaningful_keys) == 0:
        raise BreesyError(
            "No valid data found in .set file.",
            f"The file {filename} appears to be empty or corrupted. "
            "Please check that this is a valid EEGlab .set file."
        )

    if len(meaningful_keys) == 1:
        return _struct_to_dict(set_data[meaningful_keys[0]])

    eeg_like_keys = [k for k in meaningful_keys
                     if any(marker in k.lower() for marker in ['eeg', 'data', 'set'])]

    if len(eeg_like_keys) == 1:
        return _struct_to_dict(set_data[eeg_like_keys[0]])

    raise BreesyError(
        f"Could not identify EEG structure in .set file. Found multiple possible keys: {meaningful_keys}",
        "The file structure is ambiguous. Please ensure this is a valid EEGlab .set file."
    )


def _struct_to_dict(struct) -> dict:
    if isinstance(struct, dict):
        return struct

    # Numpy structured array (record array)
    if isinstance(struct, np.ndarray) and hasattr(struct, 'dtype') and struct.dtype.names:
        result = {}
        for field_name in struct.dtype.names:
            result[field_name] = struct[field_name]
        return result

    # Not a struct - return as-is wrapped in a dict
    return {'data': struct}


def _validate_continuous_data(eeg: dict, filename: str) -> None:
    # Check if this is epoched data
    trials = eeg.get('trials', 1)

    # Handle different data types that scipy.io.loadmat might return
    if isinstance(trials, np.ndarray):
        trials = int(trials.flat[0])
    else:
        trials = int(trials)

    if trials != 1:
        raise BreesyError(
            f"This EEGlab file contains epoched data ({trials} epochs/trials). "
            "Breesy currently supports only continuous EEG data.",
            "Please export your data as continuous from EEGlab before loading, "
            "or merge your epochs back into continuous data."
        )


def _extract_metadata(eeg: dict, filename: str) -> tuple[float, int, int]:
    # Extract sample rate (srate)
    if 'srate' not in eeg:
        raise BreesyError(
            "Missing 'srate' (sampling rate) field in .set file.",
            f"The file {filename} does not contain sampling rate information. "
            "Please check that this is a valid EEGlab .set file."
        )
    sample_rate = _extract_scalar_value(eeg['srate'], 'srate', filename)
    sample_rate = float(sample_rate)

    if sample_rate <= 0:
        raise BreesyError(
            f"Invalid sample rate: {sample_rate} Hz",
            "Sample rate must be positive. Please check the .set file."
        )

    # Extract number of channels (nbchan)
    if 'nbchan' not in eeg:
        raise BreesyError(
            "Missing 'nbchan' (number of channels) field in .set file.",
            f"The file {filename} does not contain channel count information. "
            "Please check that this is a valid EEGlab .set file."
        )
    nbchan = _extract_scalar_value(eeg['nbchan'], 'nbchan', filename)
    nbchan = int(nbchan)

    if nbchan <= 0:
        raise BreesyError(
            f"Invalid number of channels: {nbchan}",
            "Number of channels must be positive. Please check the .set file."
        )

    # Extract number of time points (pnts)
    if 'pnts' not in eeg:
        raise BreesyError(
            "Missing 'pnts' (number of time points) field in .set file.",
            f"The file {filename} does not contain time points information. "
            "Please check that this is a valid EEGlab .set file."
        )
    pnts = _extract_scalar_value(eeg['pnts'], 'pnts', filename)
    pnts = int(pnts)

    if pnts <= 0:
        raise BreesyError(
            f"Invalid number of time points: {pnts}",
            "Number of time points must be positive. Please check the .set file."
        )

    return sample_rate, nbchan, pnts


def _extract_channel_names(eeg: dict, nbchan: int, filename: str) -> list[str]:
    default_names = [f'Channel_{i+1}' for i in range(nbchan)]

    if 'chanlocs' not in eeg:
        return default_names

    chanlocs = eeg['chanlocs']
    if chanlocs is None or (isinstance(chanlocs, np.ndarray) and chanlocs.size == 0):
        return default_names

    try:
        labels = _extract_labels_from_chanlocs(chanlocs, nbchan)
        if labels and len(labels) == nbchan:
            return labels
    except Exception:
        pass

    return default_names


def _extract_labels_from_chanlocs(chanlocs, nbchan: int) -> list[str] | None:
    # Handle scalar numpy arrays (common with scipy.io.loadmat)
    if isinstance(chanlocs, np.ndarray) and chanlocs.shape == ():
        chanlocs = chanlocs.item()

    # Case 1: chanlocs is a structured numpy array (most common)
    if isinstance(chanlocs, np.ndarray):
        # Check if it has 'labels' field
        if hasattr(chanlocs, 'dtype') and chanlocs.dtype.names:
            if 'labels' in chanlocs.dtype.names:
                labels = chanlocs['labels']
                return _normalize_labels(labels, nbchan)

        # Check for alternative field names
        for field_name in ['label', 'name', 'channel']:
            if hasattr(chanlocs, 'dtype') and chanlocs.dtype.names and field_name in chanlocs.dtype.names:
                labels = chanlocs[field_name]
                return _normalize_labels(labels, nbchan)

    # Case 2: chanlocs is a list or dict
    if isinstance(chanlocs, (list, tuple)):
        try:
            labels = [ch.get('labels', ch.get('label', '')) if isinstance(ch, dict) else str(ch)
                     for ch in chanlocs]
            if len(labels) == nbchan:
                return [str(label).strip() for label in labels if label]
        except Exception:
            pass

    return None


def _normalize_labels(labels, nbchan: int) -> list[str] | None:
    try:
        # Convert to list if numpy array
        if isinstance(labels, np.ndarray):
            # Flatten and convert to strings
            if labels.ndim == 0:
                # Scalar - single channel
                return [str(labels.item()).strip()]
            elif labels.ndim == 1:
                # 1D array - multiple channels
                result = [str(label).strip() for label in labels]
                if len(result) == nbchan and all(result):
                    return result
            else:
                # Multi-dimensional - try to flatten
                result = [str(label).strip() for label in labels.flat]
                if len(result) == nbchan and all(result):
                    return result

        # Already a list or tuple
        elif isinstance(labels, (list, tuple)):
            result = [str(label).strip() for label in labels]
            if len(result) == nbchan and all(result):
                return result

    except Exception:
        pass

    return None


def _load_eeglab_data(eeg: dict, set_filename: str, nbchan: int, pnts: int) -> np.ndarray:
    """Load EEG data from embedded array or external .fdt file.

    EEGlab can store data in two ways:
    1. Embedded in the .set file (as a numpy array in the 'data' field)
    2. In an external binary .fdt file (referenced by filename in 'data' field)
    """
    if 'data' not in eeg:
        raise BreesyError(
            "Missing 'data' field in .set file.",
            f"The file {set_filename} does not contain EEG data. "
            "Please check that this is a valid EEGlab .set file."
        )

    data_field = eeg['data']

    # Handle scalar numpy arrays (common with scipy.io.loadmat)
    if isinstance(data_field, np.ndarray) and data_field.shape == ():
        data_field = data_field.item()

    # Case 1: Data is in an external .fdt file (referenced by filename string)
    if isinstance(data_field, str):
        fdt_filename = data_field
        return _load_fdt_data(fdt_filename, set_filename, nbchan, pnts)

    # Case 2: Data is embedded in the .set file as a numpy array
    if isinstance(data_field, np.ndarray):
        return _process_embedded_data(data_field, eeg, nbchan, pnts, set_filename)

    # Case 3: Unexpected data type
    raise BreesyInternalError(
        f"Unexpected data type in 'data' field: {type(data_field)} "
        f"The .set file {set_filename} has an unexpected structure. "
        "Expected numpy array or filename string."
    )


def _process_embedded_data(data: np.ndarray, eeg: dict, nbchan: int, pnts: int, filename: str) -> np.ndarray:
    # Validate data shape
    expected_shape = (nbchan, pnts)

    if data.shape != expected_shape:
        # Try transposed shape
        if data.shape == (pnts, nbchan):
            data = data.T
        else:
            raise BreesyInternalError(
                f"Embedded data has unexpected shape: {data.shape}. Expected {expected_shape} "
                f"The .set file {filename} has inconsistent metadata. "
                f"Expected {nbchan} channels and {pnts} time points."
            )

    # Apply calibration scaling if present
    # EEGlab sometimes stores a 'CAL' field for unit conversion (typically 1e-6 for microvolts)
    if 'CAL' in eeg:
        try:
            cal = _extract_scalar_value(eeg['CAL'], 'CAL', filename)
            data = data * float(cal)
        except Exception:
            # If calibration fails, just use raw data
            pass

    return data


def _load_fdt_data(fdt_filename: str, set_filename: str, nbchan: int, pnts: int) -> np.ndarray:
    # Resolve .fdt file path
    # The fdt_filename can be:
    # 1. Just a filename (e.g., "data.fdt") - look in same directory as .set
    # 2. A relative path (e.g., "../data/data.fdt") - resolve relative to .set
    # 3. An absolute path (e.g., "/full/path/data.fdt") - use as-is

    set_path = Path(set_filename)
    fdt_path = Path(fdt_filename)

    if not fdt_path.is_absolute():
        # Relative path or just filename - resolve relative to .set file's directory
        fdt_path = set_path.parent / fdt_path

    # If still not found, try replacing .set extension with .fdt
    if not fdt_path.exists():
        fdt_path = set_path.with_suffix('.fdt')

    # Final check - does the file exist?
    if not fdt_path.exists():
        raise BreesyError(
            f"Data file not found: {fdt_filename}",
            f"EEGlab .set files require a corresponding .fdt file. "
            f"Looked in: {fdt_path.parent}\n"
            f"Make sure the .fdt file is in the same directory as the .set file."
        )

    # Read the binary .fdt file
    return _read_fdt_file(str(fdt_path), nbchan, pnts)


def _read_fdt_file(fdt_path: str, nbchan: int, pnts: int) -> np.ndarray:
    """Read binary .fdt file and reshape to (channels, samples).

    .fdt files store data as:
    - Data type: 32-bit IEEE floating point (float32)
    - Byte order: Little-endian
    - Memory order: Fortran/column-major (samples fastest)
    """
    try:
        # Read raw binary data as float32
        data = np.fromfile(fdt_path, dtype=np.float32)

        # Validate size
        expected_size = nbchan * pnts
        if data.size != expected_size:
            raise BreesyError(
                f"Data size mismatch in .fdt file. "
                f"Expected {expected_size} values ({nbchan} channels × {pnts} samples), "
                f"but got {data.size} values.",
                "The .fdt file may be corrupted or does not match the .set metadata. "
                "Make sure the .fdt and .set files are from the same recording."
            )

        # Reshape from flat array to 2D
        # .fdt files are in Fortran (column-major) order:
        # - Samples vary fastest (time points in each channel are contiguous)
        # - Channels vary slowest
        # This means the file is organized as: [ch1_all_samples, ch2_all_samples, ...]
        # So we reshape as (pnts, nbchan) with Fortran order, then transpose
        data = data.reshape((pnts, nbchan), order='F').T

        return data

    except BreesyError:
        # Re-raise our own errors
        raise
    except PermissionError:
        raise BreesyError(
            f"Permission denied when reading .fdt file: {fdt_path}",
            "Check that you have read permissions for this file."
        )
    except Exception as e:
        raise BreesyInternalError(
            f"Failed to read .fdt file: {str(e)} "
            f"The file {fdt_path} may be corrupted or in an unexpected format."
        )


def _extract_events(eeg: dict, sample_rate: float) -> list[Event]:
    """Extract event markers from EEG structure.

    EEGlab stores events in the 'event' field as a structured array with fields like:
    - 'type': Event type/name
    - 'latency': Sample index (1-based in EEGlab)
    - 'duration': Event duration (optional)

    :return: List of Event objects sorted by index
    """
    # Check if events field exists
    if 'event' not in eeg:
        return []

    events_data = eeg['event']

    # Handle empty events
    if events_data is None:
        return []

    # Handle case where it's not an array
    if not isinstance(events_data, np.ndarray):
        return []

    # Handle empty array
    if events_data.size == 0:
        return []

    # Unwrap 0-d object arrays
    if events_data.dtype == object and events_data.shape == ():
        events_data = events_data.item()
        if not isinstance(events_data, np.ndarray):
            return []

    # Extract events from structured array
    events = []

    try:
        # Convert to 1D array if needed
        events_array = np.atleast_1d(events_data)

        for event_item in events_array:
            try:
                event_obj = _parse_single_event(event_item)
                if event_obj:
                    events.append(event_obj)
            except Exception:
                # Skip malformed events
                continue

    except Exception:
        # If anything goes wrong, return empty list
        return []

    # Sort events by index
    events.sort(key=lambda e: e.index)

    return events


def _parse_single_event(event_item) -> Event | None:
    try:
        # Extract event type/name and latency (sample index)
        event_type, latency = None, None

        if hasattr(event_item, 'dtype') and event_item.dtype.names:
            if 'type' in event_item.dtype.names:
                event_type = event_item['type']
            if 'latency' in event_item.dtype.names:
                latency = event_item['latency']
        elif isinstance(event_item, dict):
            event_type = event_item.get('type')
            latency = event_item.get('latency')

        # no latency = no sample index = no event
        if latency is None:
            return None
        if isinstance(latency, np.ndarray):
            if latency.size == 0:
                return None
            latency_value = int(latency.flat[0])
        else:
            latency_value = int(latency)
        # EEGlab uses 1-based indexing, convert to 0-based
        sample_index = latency_value - 1
        if sample_index < 0:
            return None    

        # no event type = default event
        if event_type is None:
            event_name = 'Event'
        elif isinstance(event_type, np.ndarray):
            event_name = str(event_type.flat[0]) if event_type.size > 0 else 'Event'
        elif isinstance(event_type, (bytes, np.bytes_)):
            event_name = event_type.decode('utf-8', errors='ignore')
        else:
            event_name = str(event_type)

        event_name = event_name.strip()
        if not event_name:
            event_name = 'Event'

        return Event(name=event_name, index=sample_index)

    except Exception:
        return None


def _extract_scalar_value(value, field_name: str, filename: str):
    if isinstance(value, (int, float)):
        return value
    if isinstance(value, np.ndarray):
        if value.size == 1:
            return value.flat[0]
        else:
            raise BreesyError(
                f"Expected scalar value for '{field_name}', but got array of size {value.size}",
                f"The .set file {filename} has an unexpected structure."
            )
    if isinstance(value, (list, tuple)) and len(value) == 1:
        return value[0]

    raise BreesyInternalError(
        f"Could not extract scalar value from '{field_name}' field. Got type: {type(value)} "
        f"The .set file {filename} has an unexpected structure."
    )
