from __future__ import annotations

import logging

import numpy as np

from config import MIN_ECG_SECONDS

LOGGER = logging.getLogger(__name__)


def process_ecg_signal(ecg_signal: np.ndarray, sampling_rate: int) -> dict[str, object]:
    """Clean ECG, detect R peaks, and compute RR intervals."""
    signal = np.asarray(ecg_signal, dtype=float).reshape(-1)
    min_samples = int(MIN_ECG_SECONDS * sampling_rate)
    if signal.size < min_samples:
        return _error(f"ECG signal is too short. Need at least {MIN_ECG_SECONDS} seconds.")

    if not np.isfinite(signal).all():
        signal = signal[np.isfinite(signal)]
    if signal.size < min_samples:
        return _error("ECG signal has too few finite samples after cleaning invalid values.")

    try:
        cleaned_ecg, rpeaks = _process_with_neurokit(signal, sampling_rate)
    except Exception:
        LOGGER.warning("neurokit2 ECG processing unavailable; falling back to scipy/manual peak detection.")
        cleaned_ecg, rpeaks = _process_with_fallback(signal, sampling_rate)

    rpeaks = np.asarray(rpeaks, dtype=int)
    if rpeaks.size < 3:
        return _error("Too few R peaks were detected for HR/HRV estimation.")

    rr_intervals_ms = np.diff(rpeaks) / sampling_rate * 1000.0
    if rr_intervals_ms.size == 0 or float(np.mean(rr_intervals_ms)) <= 0:
        return _error("Unable to compute valid RR intervals.")

    heart_rate = 60000.0 / float(np.mean(rr_intervals_ms))
    return {
        "ok": True,
        "heart_rate": round(heart_rate, 2),
        "rr_intervals_ms": rr_intervals_ms.round(2).tolist(),
        "rpeaks": rpeaks.tolist(),
        "cleaned_ecg": np.asarray(cleaned_ecg, dtype=float).round(6).tolist(),
        "message": "success",
    }


def _process_with_neurokit(signal: np.ndarray, sampling_rate: int) -> tuple[np.ndarray, np.ndarray]:
    import neurokit2 as nk

    cleaned = nk.ecg_clean(signal, sampling_rate=sampling_rate)
    _, info = nk.ecg_peaks(cleaned, sampling_rate=sampling_rate)
    return np.asarray(cleaned), np.asarray(info.get("ECG_R_Peaks", []), dtype=int)


def _process_with_fallback(signal: np.ndarray, sampling_rate: int) -> tuple[np.ndarray, np.ndarray]:
    centered = signal - np.median(signal)
    try:
        from scipy.signal import find_peaks

        min_distance = int(0.45 * sampling_rate)
        threshold = np.percentile(centered, 90)
        peaks, _ = find_peaks(centered, distance=min_distance, height=threshold)
        return centered, peaks
    except Exception:
        threshold = np.percentile(centered, 95)
        candidate_indices = np.where(centered >= threshold)[0]
        peaks = []
        last_peak = -int(0.45 * sampling_rate)
        for idx in candidate_indices:
            if idx - last_peak >= int(0.45 * sampling_rate):
                window_end = min(idx + int(0.2 * sampling_rate), len(centered))
                local_idx = idx + int(np.argmax(centered[idx:window_end]))
                peaks.append(local_idx)
                last_peak = local_idx
        return centered, np.asarray(peaks, dtype=int)


def _error(message: str) -> dict[str, object]:
    return {
        "ok": False,
        "heart_rate": None,
        "rr_intervals_ms": [],
        "rpeaks": [],
        "cleaned_ecg": [],
        "message": message,
    }