import importlib.util
import json
import logging
import os
import platform
import re
from enum import Enum
from pathlib import Path
from typing import Any, Callable, List, Optional

import ffmpeg
import numpy as np
from rich import print
from rich.progress import (
    Progress, SpinnerColumn, TextColumn, TimeElapsedColumn,
    BarColumn, TaskProgressColumn, MofNCompleteColumn, TimeRemainingColumn
)

# Set logging configuration
logger = logging.getLogger("little_helper")
logger.setLevel(logging.INFO)


class DeviceChoice(str, Enum):
    AUTO = "auto"
    CPU = "cpu"
    GPU = "gpu"
    MLX = "mlx"


OPTIONAL_DEVICE_DEPENDENCIES = {
    "mlx": {"extra": "mlx", "packages": [("mlx-whisper", "mlx_whisper")]}
}


def get_device(device: DeviceChoice = DeviceChoice.AUTO) -> str:
    """
    Determine the computation device based on user preference and
    availability.
    """
    torch_module = None
    torch_available = False

    if device != DeviceChoice.MLX:
        try:
            import torch as torch_module

            torch_available = True
        except ImportError:
            torch_available = False

    if device == DeviceChoice.AUTO:
        if not torch_available:
            return "cpu"
        if torch_module.cuda.is_available():
            device = "cuda:0"
        elif torch_module.backends.mps.is_available():
            device = "mlx"
        else:
            device = "cpu"
    elif device == DeviceChoice.GPU:
        if torch_available and torch_module.cuda.is_available():
            device = "cuda:0"
        else:
            print("[blue1]→ NVIDIA GPU not available or PyTorch missing. Using CPU.")
            device = "cpu"
    elif device == DeviceChoice.MLX:
        if platform.system() == "Darwin":
            device = "mlx"
        else:
            print("[blue1]→ MLX is only available on macOS. Using CPU.")
            device = "cpu"
    elif device == DeviceChoice.CPU:
        device = "cpu"
    else:
        device = "cpu"
    return device


def _normalize_device_key(device: Optional[str]) -> str:
    if not device:
        return "cpu"
    if device in ["cuda", "cuda:0", "gpu"]:
        return "cuda:0"
    if device == "mlx":
        return "mlx"
    return "cpu"


def _missing_packages(packages: list[tuple[str, str]]) -> list[str]:
    missing = []
    for pip_name, module_name in packages:
        if importlib.util.find_spec(module_name) is None:
            missing.append(pip_name)
    return missing


def check_dependencies_for_device(
    device: str, requested_device: Optional[DeviceChoice | str] = None
) -> tuple[bool, Optional[str]]:
    """
    Verify that optional dependencies for the selected device are installed.
    Returns (ok, message).
    """
    target_device = requested_device
    if isinstance(target_device, DeviceChoice):
        target_device = target_device.value
    if target_device in [None, DeviceChoice.AUTO.value]:
        target_device = device

    device_key = _normalize_device_key(target_device)
    requirements = OPTIONAL_DEVICE_DEPENDENCIES.get(device_key)
    if not requirements:
        return True, None

    missing = _missing_packages(requirements["packages"])
    if missing:
        extra = requirements.get("extra")
        extras_hint = (
            f' Install with `pip install "whispi\\[{extra}]"`.' if extra else ""
        )
        message = (
            f"Missing dependencies for device '{target_device}'. "
            f"Install the following: {', '.join(missing)}.{extras_hint}"
        )
        return False, message

    return True, None


class FilePathProcessor:
    """
    Utility class for validating various filepaths.
    """

    def __init__(self, file_formats: List[str]):
        self.file_formats = [fmt.lower() for fmt in file_formats]
        self.filepaths: List[Path] = []

    def get_filepaths(self, filepath: str):
        """
        Processes the provided filepath which can be a single file,
        a directory, or a .list file containing multiple paths. It
        validates each input and accumulates valid file paths for
        further processing.
        """
        path = Path(filepath).expanduser().resolve()

        try:
            # Handle .list file
            if path.suffix.lower() == ".list":
                if not path.is_file():
                    msg = f"The .list '{path}' does not exist or is not a file"
                    logging.error(msg)
                    print(f"→ {msg}")
                    return

                logging.info(f"Processing .list file: {path}")
                with path.open("r", encoding="utf-8") as file:
                    lpaths = set()
                    for line in file:
                        lpath = line.strip()
                        if not lpath:
                            continue
                        lpaths.add(lpath)

                    for lpath in lpaths:
                        self._process_path(lpath)
                return

            # Handle single file or directory
            else:
                self._process_path(path)

        except Exception as e:
            logging.exception(
                f"An unexpected error occurred while processing '{filepath}': {e}"
            )
            return

        # Remove duplicates by converting to a set of resolved absolute paths
        unique_filepaths = set(p.resolve() for p in self.filepaths)
        self.filepaths = list(unique_filepaths)

        # Filter out files that have already been converted
        self._filter_converted_files()

        # Final check to ensure there are files to process
        if not self.filepaths:
            logging.warning(
                f"No valid files found for processing. "
                f'Please check the provided path: "{filepath}".'
            )
            print(
                f"→ No valid files found for processing. "
                f'Please check the provided path: "{filepath}".'
            )
        else:
            logging.info(f"Total valid files to process: {len(self.filepaths)}")

    def _process_path(self, path_input: str | Path):
        """
        Processes a single path input, which can be a file or a directory.
        """
        path = Path(path_input).expanduser().resolve()

        if path.is_file():
            if path.suffix.lower() in self.file_formats:
                logging.info(f"Adding file: {path}")
                normalized_path = self._normalize_filepath(path)
                self.filepaths.append(normalized_path)
            else:
                logging.warning(
                    f"File '{path}' has unsupported format and will be skipped."
                )
                print(f"→ File '{path}' has unsupported format and will be skipped.")
        elif path.is_dir():
            logging.info(f"Processing directory: {path}")
            for file_format in self.file_formats:
                for file in path.rglob(f"*{file_format}"):
                    if file.is_file():
                        logging.debug(f"Found file: {file}")
                        normalized_path = self._normalize_filepath(file)
                        self.filepaths.append(normalized_path)
        else:
            logging.error(f'Path "{path}" does not exist or is not accessible.')
            print(f'→ Path "{path}" does not exist or is not accessible.')

    def _normalize_filepath(self, filepath: Path) -> Path:
        """
        Normalizes the filepath by replacing non-word characters with
        underscores,collapsing multiple underscores into one, and removing
        leading/trailing underscores.
        """
        new_filename = re.sub(r"\W+", "_", filepath.stem)
        new_filename = new_filename.strip("_")
        new_filename = re.sub(r"_+", "_", new_filename)

        suffix = filepath.suffix.lower()

        # Construct the new path
        new_path = filepath.parent / f"{new_filename}{suffix}"

        # Rename the file
        filepath.rename(new_path)

        return new_path.resolve()

    def _filter_converted_files(self):
        """
        Removes files that have already been converted to avoid
        redundant processing.
        """
        converted_suffix = "_converted.wav"
        original_filepaths = []
        converted_filepaths = set()

        for fp in self.filepaths:
            if fp.name.endswith(converted_suffix):
                converted_filepaths.add(fp)
            else:
                original_filepaths.append(fp)

        # Remove originals if their converted version exists
        filtered_filepaths = [
            fp
            for fp in original_filepaths
            if not (fp.with_name(fp.stem + converted_suffix) in converted_filepaths)
        ]

        # Extened filtered paths with converted paths
        filtered_filepaths.extend(converted_filepaths)

        removed_count = len(self.filepaths) - len(filtered_filepaths)
        if removed_count > 0:
            logging.info(f"Removed {removed_count} files already converted.")
        self.filepaths = filtered_filepaths


def ensure_dir(dir: Path) -> None:
    if not dir.exists():
        dir.mkdir(parents=True)
    return dir


def set_output_dir(filepath: Path, base_dir: Path) -> None:
    output_dir = base_dir / filepath.stem
    ensure_dir(output_dir)
    return output_dir


def return_valid_fileformats() -> list[str]:
    return [
        ".mp3",
        ".wav",
        ".m4a",
        ".aac",
        ".flac",
        ".ogg",
        ".mkv",
        ".mov",
        ".mp4",
        ".avi",
        ".mpeg",
        ".vob",
    ]


def load_audio_ffmpeg(filepath: str) -> np.ndarray:
    try:
        out, _ = (
            ffmpeg.input(filepath)
            .output("pipe:", format="f32le", acodec="pcm_f32le", ac=1, ar="16000")
            .run(capture_stdout=True, capture_stderr=True)
        )
    except ffmpeg.Error as e:
        raise RuntimeError(
            f"Error loading audio with ffmpeg: {e.stderr.decode()}"
        ) from e
    return np.frombuffer(out, np.float32)


def check_file_format(filepath: Path) -> tuple[Path, np.ndarray]:
    """
    Checks the format of an audio file and converts it if it doesn't meet
    specified criteria. Then, loads the audio into a 1D NumPy array.

    The function uses `ffmpeg` to probe the metadata of an audio file at
    the given `filepath`. It checks if the audio stream meets the following
    criteria:
        - Codec name: 'pcm_s16le'
        - Sample rate: 16000 Hz
        - Number of channels: 1 (mono)

    If the audio stream does not meet these criteria, the function attempts
    to convert the file to meet the required format and saves the converted
    file with a '_converted.wav' suffix in the same directory.
    Original input files are always preserved.

    Finally, it loads the audio (original or converted) as a 1D NumPy array
    and returns it.

    Args:
        filepath (Path): The path to the audio file to be checked and
            potentially converted.

    Returns:
        filepath (Path): filepath of the checked and / or converted audio file.
        np.ndarray: 1D NumPy array of the audio data.
    """
    # Define the converted file path
    new_filepath = filepath.with_name(f"{filepath.stem}_converted.wav")

    converted = False
    if new_filepath.exists():
        target_filepath = new_filepath
        converted = True
    else:
        try:
            # Probe the audio file for stream information
            probe = ffmpeg.probe(str(filepath))
            audio_streams = [
                stream for stream in probe["streams"] if stream["codec_type"] == "audio"
            ]

            if not audio_streams:
                raise ValueError(
                    f"→ No audio stream found for {filepath}. "
                    f"Please check if the file you have provided contains "
                    f"audio content."
                )

            audio_stream = audio_streams[0]
            codec_name = audio_stream.get("codec_name")
            sample_rate = int(audio_stream.get("sample_rate", 0))
            channels = int(audio_stream.get("channels", 0))

            # Check if the audio stream meets the criteria
            if codec_name != "pcm_s16le" or sample_rate != 16000 or channels != 1:
                try:
                    # Convert the file and show progress
                    run_with_progress(
                        description=(
                            f"[orchid]→ Converting file to .wav: {filepath.name}"
                        ),
                        task=lambda: convert_file_format(
                            old_filepath=filepath, new_filepath=new_filepath
                        ),
                    )
                    target_filepath = new_filepath
                    converted = True
                except Exception as e:
                    raise RuntimeError(
                        f"→ An error occurred while converting {filepath}: {e}"
                    )
            else:
                # If already in correct format, use the original file
                target_filepath = filepath

        except ffmpeg.Error as e:
            print(
                f"→ Error running ffprobe: {e} "
                "  You may have provided an unsupported file type. Please "
                "  check 'whispi --list_formats' for all supported formats."
            )

    try:
        audio_array = load_audio_ffmpeg(str(target_filepath))
    except Exception as e:
        raise RuntimeError(f"Failed to load audio from {target_filepath}: {e}") from e

    return Path(target_filepath), audio_array


def convert_file_format(old_filepath: str, new_filepath: str):
    """
    Converts a video file into an audio file in WAV format using the
    ffmpeg library.
    """
    (
        ffmpeg.input(str(old_filepath))
        .output(
            str(new_filepath),
            acodec="pcm_s16le",  # Audio codec: PCM signed 16-bit little-endian
            ar="16000",  # Sampling rate 16 KHz
            ac=1,  # Mono channel
        )
        .run(quiet=True, overwrite_output=True)
    )


def load_config(config: json) -> dict:
    with open(config, "r", encoding="utf-8") as file:
        return json.load(file)


def format_time(seconds, delimiter=",") -> str:
    """
    Function for time conversion with millisecond precision (for subtitles).
    """
    h = int(seconds // 3600)
    m = int((seconds % 3600) // 60)
    s = int(seconds % 60)
    ms = int((seconds - int(seconds)) * 1000)

    return f"{h:02}:{m:02}:{s:02}{delimiter}{ms:03}"


def format_time_simple(seconds: float) -> str:
    """
    Format time as MM:SS for simple TXT display.
    Reference: /Users/nv/proj.xt.com/demos/video2text/.claude/skills/media2text/scripts/transcribe.py
    """
    m = int(seconds // 60)
    s = int(seconds % 60)
    return f"{m:02d}:{s:02d}"


def run_with_progress(description: str, task: Callable[[], Any]) -> Any:
    """
    Helper function to run a task with a progress bar.
    """
    with Progress(
        SpinnerColumn(),
        TimeElapsedColumn(),
        TextColumn("[progress.description]{task.description}"),
    ) as progress:
        progress.add_task(description, total=None)
        return task()


def get_audio_duration(filepath: Path) -> float:
    """
    Get audio duration in seconds using ffprobe.
    """
    try:
        probe = ffmpeg.probe(str(filepath))
        duration = float(probe['format']['duration'])
        return duration
    except Exception as e:
        logger.warning(f"Could not get audio duration: {e}")
        return 0.0


def create_dual_progress() -> Progress:
    """
    Create a Progress instance for dual progress bar display.
    Used for sequential transcription and diarization.

    Includes TimeRemainingColumn to show estimated time remaining.
    """
    return Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(bar_width=30),
        TaskProgressColumn(),
        TimeElapsedColumn(),
        TimeRemainingColumn(),
    )
