"""
neuroflow.local_client -- Direct local access to NeuroFlow without a server.

Wraps the NeuroFlow DeviceManager for direct device interaction. No
running API server is needed -- everything runs in-process using
BrainFlow for hardware communication.

When the full ``neuroflow`` platform package is installed, LocalClient
also provides session management, data storage, pipeline processing,
calibration, and ML training/inference via SessionManager.

This is ideal for:
- Quick prototyping and experimentation
- Embedded applications
- Single-user desktop apps
- Testing with synthetic (simulated) devices

Usage::

    from neuroflow import LocalClient

    client = LocalClient()
    device = client.connect(board_id=-1)  # Synthetic board
    device.start_stream()

    # Get real-time features
    features = device.get_features()
    print(f"Alpha power: {features.band_powers.get('avg_alpha_power', 0)}")

    # Session management (requires neuroflow platform package)
    session_id = client.start_session(device.id, pipeline="meditation")
    client.stop_session(session_id)
    data = client.get_session_data(session_id)

    device.stop_stream()
    device.disconnect()
"""

from __future__ import annotations

import time
import uuid
from typing import Any, Optional

from neuroflow.device import LocalDevice
from neuroflow.models import Features, StreamData


class LocalClient:
    """
    Direct local client for NeuroFlow -- no server required.

    Creates and manages BrainFlow device connections directly
    in the current process. Uses the same DeviceManager that
    powers the NeuroFlow server, but without any network overhead.

    When the full neuroflow platform package is installed, also provides:
    - Session management (start/stop/list/get sessions)
    - Persistent data storage (SQLite)
    - Pipeline processing
    - Calibration with storage
    - ML model training and inference

    Requires the ``neuroflow`` package and ``brainflow`` to be installed.
    Install with: ``pip install neuroflow[local]``

    Args:
        auto_reconnect: Whether devices should automatically reconnect
            on connection loss (default True).
        buffer_size: Size of the BrainFlow ring buffer in samples
            (default 450000, roughly 30 minutes at 250 Hz).
        storage_path: SQLite database path for session storage
            (default "neuroflow_data.db").
        enable_sessions: Whether to enable session management
            (default True, falls back gracefully if platform not installed).

    Example::

        client = LocalClient()
        device = client.connect(board_id=-1)  # Synthetic
        device.start_stream()
        features = device.get_features()
        device.stop_stream()
        client.close()
    """

    def __init__(
        self,
        auto_reconnect: bool = True,
        buffer_size: int = 450000,
        storage_path: str = "neuroflow_data.db",
        enable_sessions: bool = True,
    ) -> None:
        self._auto_reconnect = auto_reconnect
        self._buffer_size = buffer_size
        self._devices: dict[str, LocalDevice] = {}
        self._session_manager: Any = None
        self._calibrations: dict[str, Any] = {}

        # Verify that neuroflow is importable
        try:
            from neuroflow.acquisition.device_manager import DeviceManager  # noqa: F401
        except ImportError:
            raise ImportError(
                "The 'neuroflow' package is required for local device access. "
                "Install it with: pip install neuroflow[local]"
            )

        # Initialize session manager if platform package is available
        if enable_sessions:
            try:
                from neuroflow.platform.session_manager import SessionManager
                self._session_manager = SessionManager()
                self._session_manager._storage_path = storage_path  # type: ignore[attr-defined]
            except ImportError:
                pass  # Platform not installed, sessions not available

    def _require_sessions(self) -> Any:
        """Ensure session manager is available, raise if not."""
        if self._session_manager is None:
            raise RuntimeError(
                "Session management requires the full neuroflow platform package. "
                "Install with: pip install neuroflow[local]"
            )
        return self._session_manager

    # ── Device Management ─────────────────────────────────────

    def connect(
        self,
        board_id: int = -1,
        params: Optional[dict[str, Any]] = None,
        timeout: float = 30.0,
    ) -> LocalDevice:
        """
        Connect to a BrainFlow-compatible device.

        Creates a DeviceManager, establishes the connection, and returns
        a LocalDevice handle for interacting with the hardware.

        Args:
            board_id: BrainFlow board identifier.
                Common values:
                - ``-1`` : Synthetic board (simulated data, no hardware)
                - ``0``  : OpenBCI Cyton
                - ``2``  : OpenBCI Ganglion
                - ``22`` : Muse S (BLE)
                - ``38`` : Muse 2 (BLE)
                See BrainFlow docs for the full list.
            params: BrainFlow connection parameters dict. Keys depend on
                the board type (e.g. ``serial_port``, ``mac_address``,
                ``ip_address``).
            timeout: Connection timeout in seconds.

        Returns:
            LocalDevice handle for the connected device.

        Raises:
            ConnectionError: If the device fails to connect.

        Example::

            # Synthetic device (no hardware needed)
            device = client.connect(board_id=-1)

            # OpenBCI Cyton via serial
            device = client.connect(
                board_id=0,
                params={"serial_port": "/dev/ttyUSB0"},
            )

            # Muse S via BLE
            device = client.connect(
                board_id=22,
                params={"mac_address": "XX:XX:XX:XX:XX:XX"},
            )
        """
        # If session manager is available, use it for device management too
        if self._session_manager is not None:
            device_id = self._session_manager.connect_device(
                board_id=board_id, params=params or {},
            )
            managed = self._session_manager.get_device(device_id)
            device = LocalDevice(
                device_manager=managed.manager,
                device_id=device_id,
                board_id=board_id,
            )
            self._devices[device_id] = device
            return device

        from neuroflow.acquisition.device_manager import DeviceManager
        from neuroflow.acquisition.config import KNOWN_DEVICES

        # Find device config in registry
        config = None
        for known in KNOWN_DEVICES:
            if known.board_id == board_id:
                config = known
                break

        # Create DeviceManager
        dm = DeviceManager(
            board_id=board_id,
            params=params or {},
            device_config=config,
            auto_reconnect=self._auto_reconnect,
            buffer_size=self._buffer_size,
        )

        # Connect
        try:
            dm.connect(timeout=timeout)
        except Exception as exc:
            raise ConnectionError(
                f"Failed to connect to device (board_id={board_id}): {exc}"
            ) from exc

        # Generate device ID
        device_id = f"local_{uuid.uuid4().hex[:8]}"

        device = LocalDevice(
            device_manager=dm,
            device_id=device_id,
            board_id=board_id,
        )
        self._devices[device_id] = device

        return device

    def list_devices(self) -> list[LocalDevice]:
        """
        List all currently connected local devices.

        Returns:
            List of LocalDevice handles.
        """
        return list(self._devices.values())

    def get_device(self, device_id: str) -> LocalDevice:
        """
        Get a connected device by its identifier.

        Args:
            device_id: Device identifier returned by connect().

        Returns:
            LocalDevice handle.

        Raises:
            KeyError: If the device is not found.
        """
        if device_id not in self._devices:
            raise KeyError(f"Device '{device_id}' not found. Connected: {list(self._devices.keys())}")
        return self._devices[device_id]

    def disconnect(self, device_id: str) -> None:
        """
        Disconnect a specific device.

        Args:
            device_id: Device identifier to disconnect.

        Raises:
            KeyError: If the device is not found.
        """
        device = self.get_device(device_id)
        if device.is_streaming:
            device.stop_stream()
        device.disconnect()
        del self._devices[device_id]
        if self._session_manager is not None:
            try:
                self._session_manager.disconnect_device(device_id)
            except Exception:
                pass

    def disconnect_all(self) -> None:
        """Disconnect all connected devices and release resources."""
        for device_id in list(self._devices.keys()):
            try:
                self.disconnect(device_id)
            except Exception:
                pass

    # ── Session Management ────────────────────────────────────

    def start_session(
        self,
        device_id: str,
        name: str = "",
        pipeline: str = "",
        subject_id: str = "",
        notes: str = "",
        tags: Optional[list[str]] = None,
    ) -> str:
        """
        Start a recording session.

        Requires the full neuroflow platform package.

        Args:
            device_id: ID of a connected device.
            name: Session name.
            pipeline: Processing pipeline preset name.
            subject_id: Optional test subject identifier.
            notes: Free-form notes.
            tags: Categorization tags.

        Returns:
            session_id string.
        """
        sm = self._require_sessions()
        return sm.start_session(
            device_id=device_id,
            name=name,
            pipeline_name=pipeline,
            subject_id=subject_id,
            notes=notes,
            tags=tags or [],
        )

    def stop_session(self, session_id: str) -> dict:
        """Stop a session and return info."""
        sm = self._require_sessions()
        info = sm.stop_session(session_id)
        return info.model_dump()

    def get_session(self, session_id: str) -> dict:
        """Get session info."""
        sm = self._require_sessions()
        return sm.get_session(session_id).model_dump()

    def list_sessions(self, limit: int = 50) -> list[dict]:
        """List sessions."""
        sm = self._require_sessions()
        return [s.model_dump() for s in sm.list_sessions(limit)]

    def get_session_data(self, session_id: str) -> dict:
        """
        Get session data.

        Returns dict with:
        - eeg_data: numpy array [channels x samples]
        - timestamps: numpy array
        - features: dict
        - sampling_rate: int
        """
        sm = self._require_sessions()
        return sm.get_session_data(session_id)

    def get_features(self, session_id: str) -> dict:
        """Get computed features for a session."""
        sm = self._require_sessions()
        return sm.storage.load_features(session_id)

    def delete_session(self, session_id: str) -> None:
        """Delete a session."""
        sm = self._require_sessions()
        sm.delete_session(session_id)

    # ── Event Markers ─────────────────────────────────────────

    def add_marker(
        self,
        session_id: str,
        label: str,
        timestamp: Optional[float] = None,
        metadata: Optional[dict] = None,
    ) -> dict:
        """Add an event marker to a session."""
        sm = self._require_sessions()
        import secrets as _secrets
        marker = {
            "marker_id": f"mrk_{_secrets.token_hex(6)}",
            "session_id": session_id,
            "label": label,
            "timestamp": timestamp or time.time(),
            "metadata": metadata or {},
        }
        sm.storage.store_marker(session_id, marker)
        return marker

    def get_markers(self, session_id: str) -> list[dict]:
        """Get all markers for a session (sorted by timestamp)."""
        sm = self._require_sessions()
        return sm.storage.load_markers(session_id)

    # ── Data Export ────────────────────────────────────────────

    def export_session(self, session_id: str, format: str = "json") -> Any:
        """
        Export session data.

        Args:
            session_id: Session ID.
            format: 'csv' or 'json'.

        Returns:
            dict (json) or str (csv).
        """
        sm = self._require_sessions()
        data = sm.get_session_data(session_id)
        if format == "json":
            return {
                "session_id": session_id,
                "sampling_rate": data.get("sampling_rate", 250),
                "num_channels": data.get("num_channels", 0),
                "num_samples": data.get("num_samples", 0),
                "eeg_data": data["eeg_data"].tolist() if hasattr(data["eeg_data"], "tolist") else data["eeg_data"],
                "timestamps": data["timestamps"].tolist() if hasattr(data.get("timestamps", []), "tolist") else data.get("timestamps", []),
                "features": data.get("features", {}),
            }
        elif format == "csv":
            eeg = data["eeg_data"]
            lines = []
            n_ch = eeg.shape[0] if hasattr(eeg, "shape") and len(eeg.shape) > 1 else 1
            header = ",".join([f"CH{i}" for i in range(n_ch)])
            lines.append(header)
            if hasattr(eeg, "shape") and len(eeg.shape) == 2:
                for s in range(eeg.shape[1]):
                    lines.append(",".join(str(eeg[ch, s]) for ch in range(n_ch)))
            return "\n".join(lines)
        else:
            raise ValueError(f"Local export for '{format}' not supported. Use 'json' or 'csv'.")

    # ── Pipeline Management ───────────────────────────────────

    def list_pipelines(self) -> list[dict]:
        """List available pipelines."""
        sm = self._require_sessions()
        return sm.list_pipelines()

    def get_pipeline(self, name: str) -> dict:
        """Get pipeline configuration."""
        sm = self._require_sessions()
        p = sm.get_pipeline(name)
        return p.to_dict()

    def process(
        self,
        data: Any = None,
        session_id: str = "",
        pipeline: str = "eeg_standard",
        sampling_rate: int = 250,
    ) -> Optional[dict]:
        """
        Process data with a pipeline (one-shot).

        Args:
            data: Numpy array [channels x samples].
            session_id: Session ID (alternative to raw data).
            pipeline: Pipeline name.
            sampling_rate: Sampling rate in Hz.
        """
        sm = self._require_sessions()
        if data is not None:
            return sm.process_data(data, pipeline, sampling_rate)
        elif session_id:
            sess_data = sm.get_session_data(session_id)
            return sm.process_data(sess_data["eeg_data"], pipeline, sampling_rate)
        return None

    # ── Calibration ────────────────────────────────────────────

    def start_calibration(
        self,
        device_id: str,
        user_id: str,
        rest_duration_seconds: float = 120.0,
        focus_duration_seconds: float = 120.0,
        min_quality_score: float = 0.3,
    ) -> dict:
        """
        Start a calibration session.

        Args:
            device_id: ID of the streaming device.
            user_id: User to calibrate for.
            rest_duration_seconds: Rest phase duration (sec).
            focus_duration_seconds: Focus phase duration (sec).
            min_quality_score: Minimum signal quality (0.0-1.0).

        Returns:
            dict with calibration_id and progress.
        """
        sm = self._require_sessions()
        from neuroflow.calibration import CalibrationProtocol, CalibrationSettings

        managed = sm.get_device(device_id)
        if not managed.manager.is_streaming:
            raise RuntimeError("Device must be streaming for calibration")

        settings = CalibrationSettings(
            rest_duration_seconds=rest_duration_seconds,
            focus_duration_seconds=focus_duration_seconds,
            min_quality_score=min_quality_score,
        )

        pipeline = None
        try:
            pipeline = sm.get_pipeline("eeg_standard")
        except Exception:
            pass

        protocol = CalibrationProtocol(
            device_manager=managed.manager,
            pipeline=pipeline,
            settings=settings,
            sampling_rate=managed.manager.sampling_rate,
        )
        protocol.start_phase("rest")

        import secrets as _secrets
        cal_id = f"cal_{_secrets.token_hex(8)}"

        self._calibrations[cal_id] = {
            "protocol": protocol,
            "device_id": device_id,
            "user_id": user_id,
            "calibration_id": cal_id,
        }

        progress = protocol.progress
        progress["calibration_id"] = cal_id
        return progress

    def calibration_collect(self, calibration_id: str) -> dict:
        """Collect a data window for the current calibration phase."""
        cal = self._calibrations.get(calibration_id)
        if not cal:
            raise KeyError(f"Calibration '{calibration_id}' not found")
        cal["protocol"].collect_window(window_seconds=2.0)
        return cal["protocol"].progress

    def calibration_switch_phase(self, calibration_id: str, phase: str) -> dict:
        """Switch to a different calibration phase ('rest' or 'focus')."""
        if phase not in ("rest", "focus"):
            raise ValueError("Phase must be 'rest' or 'focus'")

        cal = self._calibrations.get(calibration_id)
        if not cal:
            raise KeyError(f"Calibration '{calibration_id}' not found")
        cal["protocol"].stop_phase()
        cal["protocol"].start_phase(phase)
        return cal["protocol"].progress

    def complete_calibration(self, calibration_id: str) -> dict:
        """Complete calibration and compute baseline profile."""
        sm = self._require_sessions()
        cal = self._calibrations.get(calibration_id)
        if not cal:
            raise KeyError(f"Calibration '{calibration_id}' not found")

        cal["protocol"].stop_phase()
        profile = cal["protocol"].complete(
            user_id=cal["user_id"],
            device_id=cal["device_id"],
        )

        profile_dict = profile.to_dict()
        sm.storage.store_calibration(profile_dict)
        self._calibrations.pop(calibration_id, None)
        return profile_dict

    def get_calibration_status(self, calibration_id: str) -> dict:
        """Get progress of an active calibration."""
        cal = self._calibrations.get(calibration_id)
        if not cal:
            raise KeyError(f"Calibration '{calibration_id}' not found")
        return cal["protocol"].progress

    def get_user_calibration(self, user_id: str) -> Optional[dict]:
        """Get the most recent calibration profile for a user."""
        sm = self._require_sessions()
        return sm.storage.load_user_calibration(user_id)

    def get_calibration_profile(self, calibration_id: str) -> Optional[dict]:
        """Get a specific calibration profile."""
        sm = self._require_sessions()
        return sm.storage.load_calibration(calibration_id)

    def delete_calibration(self, calibration_id: str) -> None:
        """Delete a calibration profile."""
        sm = self._require_sessions()
        sm.storage.delete_calibration(calibration_id)

    # ── Models & Training ────────────────────────────────────

    def _get_registry(self) -> Any:
        """Lazy-init model registry."""
        if not hasattr(self, "_registry") or self._registry is None:
            from neuroflow.ai.registry import ModelRegistry
            models_dir = getattr(self, "_models_dir", "./neuroflow_models")
            self._registry = ModelRegistry(models_dir=models_dir)
        return self._registry

    def list_models(self, status: Optional[str] = None) -> list[dict]:
        """List registered models."""
        registry = self._get_registry()
        models = registry.list_models(status=status)
        return [
            m.metadata.to_dict() if hasattr(m, "metadata") else {"name": str(m)}
            for m in models
        ]

    def load_model(self, name: str, version: Optional[str] = None) -> dict:
        """Load a model from the registry."""
        registry = self._get_registry()
        registered = registry.load(name, version=version)
        return registered.metadata.to_dict() if hasattr(registered, "metadata") else {}

    def predict(self, model_name: str, features: Any,
                session_id: Optional[str] = None) -> dict:
        """
        Run inference with a registered model.

        Args:
            model_name: Model name in the registry.
            features: Feature dict or numpy array.
            session_id: Optional session to extract features from.
        """
        import numpy as np

        registry = self._get_registry()
        registered = registry.load(model_name)
        model = registered.model

        if isinstance(features, dict):
            feature_values = np.array([list(features.values())])
        elif isinstance(features, np.ndarray):
            feature_values = features if features.ndim == 2 else features.reshape(1, -1)
        else:
            feature_values = np.array(features).reshape(1, -1)

        label = int(model.predict(feature_values)[0])
        probas = {}
        if hasattr(model, "predict_proba"):
            proba = model.predict_proba(feature_values)[0]
            class_names = (
                registered.metadata.class_names
                if hasattr(registered, "metadata") and registered.metadata.class_names
                else [str(i) for i in range(len(proba))]
            )
            probas = {name: float(p) for name, p in zip(class_names, proba)}

        return {
            "label": label,
            "confidence": max(probas.values()) if probas else 1.0,
            "probabilities": probas,
            "model": model_name,
        }

    def quick_train(
        self,
        X: Any,
        y: Any,
        class_names: Optional[list[str]] = None,
        architecture: str = "auto",
        name: str = "model",
    ) -> dict:
        """
        Train a model on feature data.

        Args:
            X: Feature matrix [n_samples x n_features].
            y: Labels [n_samples].
            class_names: Class names.
            architecture: Model architecture (or "auto").
            name: Model name.

        Returns:
            Dict with training results (accuracy, f1, etc.)
        """
        from neuroflow.ai.training import quick_train as _quick_train

        registry = self._get_registry()
        result = _quick_train(
            X=X, y=y,
            class_names=class_names,
            architecture=architecture,
            name=name,
            registry=registry,
        )
        return {
            "success": result.success,
            "error": result.error,
            "metrics": result.metrics.to_dict() if result.success else {},
            "model_name": name,
        }

    # ── System ────────────────────────────────────────────────

    def status(self) -> dict:
        """System status."""
        if self._session_manager is not None:
            return self._session_manager.get_status()
        return {
            "devices": len(self._devices),
            "mode": "local",
            "sessions_available": False,
        }

    # ── Lifecycle ─────────────────────────────────────────────

    def close(self) -> None:
        """Disconnect all devices and clean up."""
        self.disconnect_all()
        if self._session_manager is not None:
            try:
                self._session_manager.shutdown()
            except Exception:
                pass

    @staticmethod
    def list_supported_boards() -> dict[str, int]:
        """
        List commonly used BrainFlow board IDs.

        Returns:
            Dictionary mapping board name to board_id.

        Example::

            boards = LocalClient.list_supported_boards()
            # {"synthetic": -1, "cyton": 0, "ganglion": 2, ...}
        """
        return {
            "synthetic": -1,
            "cyton": 0,
            "ganglion": 2,
            "cyton_daisy": 6,
            "muse_s_bled": 22,
            "muse_2_bled": 38,
            "brainbit": 7,
            "notion_1": 13,
            "notion_2": 14,
            "crown": 23,
            "enophone": 37,
            "ant_neuro_ee_410": 24,
            "ant_neuro_ee_411": 25,
            "ant_neuro_ee_430": 28,
        }

    def __enter__(self) -> "LocalClient":
        return self

    def __exit__(self, *args: Any) -> None:
        self.close()

    def __repr__(self) -> str:
        count = len(self._devices)
        devices_str = f"{count} device{'s' if count != 1 else ''}"
        sessions = " +sessions" if self._session_manager is not None else ""
        return f"neuroflow.LocalClient({devices_str} connected{sessions})"
