import io
import logging
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any, BinaryIO, TypedDict, cast

import numpy as np
import pandas as pd
import soundfile as sf
from numpy.typing import ArrayLike
from scipy.io import wavfile

logger = logging.getLogger(__name__)


class CSVFileInfoParams(TypedDict, total=False):
    """Type definition for CSV file reader parameters in get_file_info.

    Parameters
    ----------
    delimiter : str
        Delimiter character. Default is ",".
    header : Optional[int]
        Row number to use as header. Default is 0 (first row).
        Set to None if no header.
    time_column : Union[int, str]
        Index or name of the time column. Default is 0.
    """

    delimiter: str
    header: int | None
    time_column: int | str


class CSVGetDataParams(TypedDict, total=False):
    """Type definition for CSV file reader parameters in get_data.

    Parameters
    ----------
    delimiter : str
        Delimiter character. Default is ",".
    header : Optional[int]
        Row number to use as header. Default is 0.
    time_column : Union[int, str]
        Index or name of the time column. Default is 0.
    """

    delimiter: str
    header: int | None
    time_column: int | str


class FileReader(ABC):
    """Base class for audio file readers."""

    # Class attribute for supported file extensions
    supported_extensions: list[str] = []

    @classmethod
    @abstractmethod
    def get_file_info(
        cls,
        path: str | Path | bytes | bytearray | memoryview | BinaryIO,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Get basic information about the audio file.

        Args:
            path: Path to the file.
            **kwargs: Additional parameters specific to the file reader.

        Returns:
            Dictionary containing file information including:
            - samplerate: Sampling rate in Hz
            - channels: Number of channels
            - frames: Total number of frames
            - format: File format
            - duration: Duration in seconds
        """
        pass  # pragma: no cover

    @classmethod
    @abstractmethod
    def get_data(
        cls,
        path: str | Path | bytes | bytearray | memoryview | BinaryIO,
        channels: list[int],
        start_idx: int,
        frames: int,
        **kwargs: Any,
    ) -> ArrayLike:
        """Read audio data from the file.

        Args:
            path: Path to the file.
            channels: List of channel indices to read.
            start_idx: Starting frame index.
            frames: Number of frames to read.
            **kwargs: Additional parameters specific to the file reader.

        Returns:
            Array of shape (channels, frames) containing the audio data.
        """
        pass  # pragma: no cover

    @classmethod
    def can_read(cls, path: str | Path) -> bool:
        """Check if this reader can handle the file based on extension."""
        ext = Path(path).suffix.lower()
        return ext in cls.supported_extensions


class SoundFileReader(FileReader):
    """Audio file reader using SoundFile library."""

    # SoundFile supported formats
    supported_extensions = [".wav", ".flac", ".ogg", ".aiff", ".aif", ".snd"]

    @classmethod
    def get_file_info(
        cls,
        path: str | Path | bytes | bytearray | memoryview | BinaryIO,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Get basic information about the audio file."""
        info = sf.info(_prepare_file_source(path))
        return {
            "samplerate": info.samplerate,
            "channels": info.channels,
            "frames": info.frames,
            "format": info.format,
            "subtype": info.subtype,
            "duration": info.frames / info.samplerate,
        }

    @classmethod
    def get_data(
        cls,
        path: str | Path | bytes | bytearray | memoryview | BinaryIO,
        channels: list[int],
        start_idx: int,
        frames: int,
        normalize: bool = False,
        **kwargs: Any,
    ) -> ArrayLike:
        """Read audio data from the file.

        Args:
            normalize: When False (default) and the source is a WAV file path,
                return raw integer PCM samples cast to float32 via
                scipy.io.wavfile.read. For non-WAV formats or in-memory sources,
                always uses soundfile (returning float32 normalized to [-1.0, 1.0]).
                When True, return float32 data normalized to [-1.0, 1.0] via soundfile.
        """
        logger.debug(f"Reading {frames} frames from {path!r} starting at {start_idx}")

        is_wav = isinstance(path, (str, Path)) and Path(path).suffix.lower() == ".wav"
        if not normalize and is_wav:
            # Use scipy to return raw integer samples (no normalization), cast to float32.
            source = _prepare_file_source(path)
            _sr, raw = wavfile.read(source)
            if raw.ndim == 1:
                raw = np.expand_dims(raw, axis=0)  # mono: (1, samples)
            else:
                raw = raw.T  # stereo: (channels, samples)

            # Only reindex channels when the requested selection is not the identity.
            if channels != list(range(raw.shape[0])):
                raw = raw[channels]

            result: ArrayLike = raw[:, start_idx : start_idx + frames].astype(
                np.float32,
                copy=False,
            )
            if not isinstance(result, np.ndarray):
                raise ValueError("Unexpected data type after reading file")
            logger.debug(f"File read complete (raw), returning data with shape {result.shape}")
            return result

        with sf.SoundFile(_prepare_file_source(path)) as f:
            if start_idx > 0:
                f.seek(start_idx)
            data = f.read(frames=frames, dtype="float32", always_2d=True)

            # Select requested channels
            data = data[:, channels]

            # Transpose to get (channels, samples) format
            result = data.T
            if not isinstance(result, np.ndarray):
                raise ValueError("Unexpected data type after reading file")

        _shape = result.shape
        logger.debug(f"File read complete, returning data with shape {_shape}")
        return result


class CSVFileReader(FileReader):
    """CSV file reader for time series data."""

    # CSV supported formats
    supported_extensions = [".csv"]

    @classmethod
    def get_file_info(
        cls,
        path: str | Path | bytes | bytearray | memoryview | BinaryIO,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Get basic information about the CSV file.

        Parameters
        ----------
        path : Union[str, Path]
            Path to the CSV file.
        **kwargs : Any
            Additional parameters for CSV reading. Supported parameters:

            - delimiter : str, default=","
                Delimiter character.
            - header : Optional[int], default=0
                Row number to use as header. Set to None if no header.
            - time_column : Union[int, str], default=0
                Index or name of the time column.

        Returns
        -------
        dict[str, Any]
            Dictionary containing file information including:
            - samplerate: Estimated sampling rate in Hz
            - channels: Number of data channels (excluding time column)
            - frames: Total number of frames
            - format: "CSV"
            - duration: Duration in seconds (or None if cannot be calculated)
            - ch_labels: List of channel labels

        Notes
        -----
        This method accepts CSV-specific parameters through kwargs.
        See CSVFileInfoParams for supported parameter types.
        """
        # Extract parameters with defaults
        delimiter: str = kwargs.get("delimiter", ",")
        header: int | None = kwargs.get("header", 0)
        time_column: int | str = kwargs.get("time_column", 0)

        # Read first few lines to determine structure
        df = pd.read_csv(_prepare_file_source(path), delimiter=delimiter, header=header)

        # Estimate sampling rate from first column (assuming it's time)
        try:
            # Get time column as Series
            if isinstance(time_column, str):
                time_series = df[time_column]
            else:
                time_series = df.iloc[:, time_column]
            time_values = np.array(time_series.values)
            if len(time_values) > 1:
                # Use round() instead of int() to handle floating-point precision issues
                estimated_sr = round(1 / np.mean(np.diff(time_values)))
            else:
                estimated_sr = 0  # Cannot determine from single row
        except Exception:
            estimated_sr = 0  # Default if can't calculate

        frames = df.shape[0]
        duration = frames / estimated_sr if estimated_sr > 0 else None

        # Return file info
        return {
            "samplerate": estimated_sr,
            "channels": df.shape[1] - 1,  # Assuming first column is time
            "frames": frames,
            "format": "CSV",
            "duration": duration,
            "ch_labels": df.columns[1:].tolist(),  # Assuming first column is time
        }

    @classmethod
    def get_data(
        cls,
        path: str | Path | bytes | bytearray | memoryview | BinaryIO,
        channels: list[int],
        start_idx: int,
        frames: int,
        **kwargs: Any,
    ) -> ArrayLike:
        """Read data from the CSV file.

        Parameters
        ----------
        path : Union[str, Path]
            Path to the CSV file.
        channels : list[int]
            List of channel indices to read.
        start_idx : int
            Starting frame index.
        frames : int
            Number of frames to read.
        **kwargs : Any
            Additional parameters for CSV reading. Supported parameters:

            - delimiter : str, default=","
                Delimiter character.
            - header : Optional[int], default=0
                Row number to use as header.
            - time_column : Union[int, str], default=0
                Index or name of the time column.

        Returns
        -------
        ArrayLike
            Array of shape (channels, frames) containing the data.

        Notes
        -----
        This method accepts CSV-specific parameters through kwargs.
        See CSVGetDataParams for supported parameter types.
        """
        # Extract parameters with defaults
        time_column: int | str = kwargs.get("time_column", 0)
        delimiter: str = kwargs.get("delimiter", ",")
        header: int | None = kwargs.get("header", 0)

        logger.debug(f"Reading CSV data from {path!r} starting at {start_idx}")

        # Read the CSV file
        df = pd.read_csv(_prepare_file_source(path), delimiter=delimiter, header=header)

        # Remove time column
        df = df.drop(columns=[time_column] if isinstance(time_column, str) else df.columns[time_column])

        # Select requested channels - adjust indices to account for time column removal
        if channels:
            try:
                data_df = df.iloc[:, channels]
            except IndexError:
                raise ValueError(f"Requested channels {channels} out of range")
        else:
            data_df = df

        # Handle start_idx and frames for partial reading
        end_idx = start_idx + frames if frames > 0 else None
        data_df = data_df.iloc[start_idx:end_idx]

        # Convert to numpy array and transpose to (channels, samples) format
        result = data_df.values.T

        if not isinstance(result, np.ndarray):
            raise ValueError("Unexpected data type after reading file")

        _shape = result.shape
        logger.debug(f"CSV read complete, returning data with shape {_shape}")
        return result


# Registry of available file readers
_file_readers = [SoundFileReader(), CSVFileReader()]


def _normalize_extension(file_type: str | None) -> str | None:
    if not file_type:
        return None
    ext = file_type.lower()
    if not ext.startswith("."):
        ext = f".{ext}"
    return ext


def _prepare_file_source(
    source: str | Path | bytes | bytearray | memoryview | BinaryIO,
) -> str | BinaryIO:
    if isinstance(source, (bytes, bytearray, memoryview)):
        return io.BytesIO(bytes(source))
    if hasattr(source, "read"):
        file_obj = cast(BinaryIO, source)
        try:
            file_obj.seek(0)
        except Exception:
            # Some file-like objects are not seekable or may reject seek(0).
            # In that case, continue using the current position without failing.
            logger.debug(
                "Could not seek to start of file-like object; continuing from current position",
                exc_info=True,
            )
        return file_obj
    return str(source)


def get_file_reader(
    path: str | Path | bytes | bytearray | memoryview | BinaryIO,
    *,
    file_type: str | None = None,
) -> FileReader:
    """Get an appropriate file reader for the given path or file type."""
    path_str = str(path)
    ext = _normalize_extension(file_type)
    if ext is None and isinstance(path, (str, Path)):
        ext = Path(path).suffix.lower()
    if not ext:
        raise ValueError(
            "File type is required when the extension is missing\n"
            "  Cannot determine format without an extension\n"
            "  Provide file_type like '.wav' or '.csv'"
        )

    # Try each reader in order
    for reader in _file_readers:
        if ext in reader.__class__.supported_extensions:
            logger.debug(f"Using {reader.__class__.__name__} for {path_str}")
            return reader

    # If no reader found, raise error
    raise ValueError(f"No suitable file reader found for {path_str}")


def register_file_reader(reader_class: type) -> None:
    """Register a new file reader."""
    reader = reader_class()
    _file_readers.append(reader)
    logger.debug(f"Registered new file reader: {reader_class.__name__}")
