# (c) 2026 Applied Brain Research
#
# All information contained herein is and remains the property of Applied Brain
# Research. The intellectual and technical concepts contained herein are
# proprietary to Applied Brain Research and may be covered by U.S. and Foreign
# Patents, patents in process, and are protected by trade secret or copyright
# law. Dissemination of this information or reproduction of this material is
# strictly forbidden unless prior written permission is obtained from Applied
# Brain Research. Access to the source code contained herein is hereby
# forbidden to anyone except current Applied Brain Research employees,
# contractors or other outside parties that have executed Confidentiality
# and/or Non-disclosure agreements explicitly covering such access.
#
# The copyright notice above does not evidence any actual or intended
# publication or disclosure of this source code, which includes information
# that is confidential and/or proprietary, and is a trade secret, of Applied
# Brain Research. ANY REPRODUCTION, MODIFICATION, DISTRIBUTION, PUBLIC
# PERFORMANCE, OR PUBLIC DISPLAY OF OR THROUGH USE OF THIS SOURCE CODE WITHOUT
# THE EXPRESS WRITTEN CONSENT OF APPLIED BRAIN RESEARCH IS STRICTLY PROHIBITED,
# AND IN VIOLATION OF APPLICABLE LAWS AND INTERNATIONAL TREATIES. THE RECEIPT
# OR POSSESSION OF THIS SOURCE CODE AND/OR RELATED INFORMATION DOES NOT CONVEY
# OR IMPLY ANY RIGHTS TO REPRODUCE, DISCLOSE OR DISTRIBUTE ITS CONTENTS, OR TO
# MANUFACTURE, USE, OR SELL ANYTHING THAT IT MAY DESCRIBE, IN WHOLE OR IN PART.

"""
Automatic speech recognition (ASR) wrapper and chunk parser.

Author: Andreas Stöckel (Applied Brain Research)
Author: Pawel Jaworski (Applied Brain Research)
"""

from __future__ import annotations

import contextlib
import ctypes
import enum
import logging
import time
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from types import TracebackType

from abr_sdk import cabi
from abr_sdk.cabi import AsrTextChunkType  # deliberate re-export
from abr_sdk.core import Application, Buffer, EventFlags, EventSet, Library, VariantDict
from abr_sdk.exceptions import AbrSdkError, AbrSdkUnexpectedState


class AsrMode(enum.Enum):
    """Decoder mode selecting the latency / accuracy trade-off."""

    FAST = "fast"
    ACCURATE = "accurate"

    # When we bump the min Python to 3.11+, drop this and inherit from
    # `enum.StrEnum` instead.
    def __str__(self) -> str:
        return self.value


class Asr(Application):
    """
    An ABR instance with automatic speech recognition support.

    Can be constructed from a library path with keyword arguments:

        with Asr("libabr-asr.so") as asr:

    **Simple (blocking) API** -- process an entire audio clip at once::

        transcript = asr.process(pcm_bytes)
        print(transcript.text)

    **Streaming API** -- push audio incrementally::

        transcript = AsrTranscript()
        asr.push(chunk1, on_chunk=transcript.chunks.append)
        asr.push(chunk2, on_chunk=transcript.chunks.append)
        asr.wait_for_completion()
        print(transcript.text)

    For finer control over the streaming event loop, use :class:`Processor`
    directly.
    """

    def __init__(
        self,
        lib_or_path: str | Path | Library,
        *,
        mode: AsrMode | None = None,
        enable_spellcheck: bool | None = None,
        enable_punctuation: bool | None = None,
        lib_search_paths: list[str | Path] | None = None,
        use_default_lib_search_paths: bool = True,
        resources_dir: str | Path | None = None,
        logger: logging.Logger | None = None,
    ) -> None:
        # ASR-specific config. Keys the caller did not provide are
        # omitted so the loaded library picks its own default.
        config: VariantDict = {}
        if mode is not None:
            config["asr_mode"] = mode.value
        if enable_spellcheck is not None:
            config["asr_enable_spellcheck"] = enable_spellcheck
        if enable_punctuation is not None:
            config["asr_enable_pcr"] = enable_punctuation

        super().__init__(
            lib_or_path=lib_or_path,
            lib_search_paths=lib_search_paths,
            use_default_lib_search_paths=use_default_lib_search_paths,
            app_type="asr",
            resources_dir=resources_dir,
            config=config,
            logger=logger,
        )

        # Free the native handle if the ASR-specific initialisation fails;
        # __exit__ never runs when __init__ raises.
        with contextlib.ExitStack() as close_on_error:
            close_on_error.callback(self.close)
            if not self._cabi.abr_asr_available(self._handle):
                raise AbrSdkError("The loaded application does not support ASR")
            self.input_buffer = Buffer(
                self._cabi, self._cabi.abr_asr_get_input_buffer(self._handle)
            )
            self.text_chunk_output_buffer = Buffer(
                self._cabi,
                self._cabi.abr_asr_get_text_chunk_output_buffer(self._handle),
            )
            close_on_error.pop_all()
        self._processor: Processor | None = None

    def __enter__(self) -> Asr:
        return self

    def flush(self) -> None:
        """Flush the ASR pipeline to finish processing remaining audio."""
        self._ensure_alive()
        self._cabi.abr_asr_flush(self._handle)

    def process(self, data: bytes) -> AsrTranscript:
        """
        Process PCM audio data and return the complete transcript.

        This is a synchronous/blocking call that pushes all *data* through
        the ASR pipeline, waits for the neural network to finish, and
        returns a :class:`AsrTranscript` containing the result. Cannot be
        used while a streaming session started with :meth:`push` is in
        progress.

        Parameters
        ----------
        data
            PCM audio as a little-endian 16-bit byte array.
        """
        self._ensure_alive()
        if self._processor is not None:
            raise AbrSdkError(
                "A streaming session is in progress; call wait_for_completion() "
                "before using process()"
            )
        transcript = AsrTranscript()
        with Processor(self, on_chunk=transcript.chunks.append) as proc:
            proc.process_and_wait_for_output(data, 0, True)
        return transcript

    def push(
        self,
        data: bytes,
        *,
        on_chunk: Callable[[AsrChunk], None] | None = None,
        output_poll_timeout_ms: int = 0,
    ) -> None:
        """
        Push PCM audio data into the ASR network (streaming API).

        On the first call an internal :class:`Processor` is created with
        *on_chunk* as the listener callback.  Subsequent calls reuse the
        same processor (the *on_chunk* argument is ignored after the first
        call).  Call :meth:`wait_for_completion` after the last audio
        chunk has been pushed.

        Parameters
        ----------
        data
            PCM audio as little-endian 16-bit bytes.
        on_chunk
            Callback invoked for each transcribed text chunk.  Only used
            on the first call (when the internal processor is created).
        output_poll_timeout_ms
            Extra time in milliseconds to spend waiting for output after
            the input has been pushed.  ``0`` (the default) returns as
            soon as all input bytes have been consumed.
        """
        self._ensure_alive()
        if output_poll_timeout_ms < 0:
            raise ValueError("output_poll_timeout_ms must not be negative")
        if self._processor is None:
            self._processor = Processor(self, on_chunk=on_chunk)
        self._processor.process_and_wait_for_output(data, output_poll_timeout_ms, False)

    def wait_for_completion(self) -> None:
        """
        Block until all previously pushed data is fully processed.

        The ``on_chunk`` callback may be invoked during this call.  When
        this method returns, the internal processor is closed and a new
        streaming session can be started by calling :meth:`push` again.
        """
        self._ensure_alive()
        if self._processor is None:
            return
        try:
            self._processor.process_and_wait_for_output(None, 0, True)
        finally:
            self._processor.close()
            self._processor = None

    def close(self) -> None:
        """Release all resources held by this instance."""
        if self._closed:
            return
        proc = getattr(self, "_processor", None)
        if proc is not None:
            proc.close()
            self._processor = None
        super().close()


@dataclass
class AsrChunk:
    """
    A text chunk produced by the ASR subsystem.

    Parse from raw buffer output with :meth:`parse`.  Apply to a running
    transcript with :meth:`update`.
    """

    SIZE = ctypes.sizeof(cabi.AsrTextChunk)

    type: AsrTextChunkType
    replace_byte_offset_begin: int
    replace_byte_offset_end: int
    data: bytes

    @staticmethod
    def parse(raw: bytes | bytearray) -> AsrChunk:
        """
        Parse *raw* bytes from the ASR output buffer into an :class:`AsrChunk`.

        *raw* must be exactly :attr:`SIZE` bytes.
        """
        if len(raw) != AsrChunk.SIZE:
            raise AbrSdkUnexpectedState(
                f"Expected {AsrChunk.SIZE} bytes, got {len(raw)}"
            )
        c = cabi.AsrTextChunk.from_buffer_copy(raw)
        if c.replace_byte_offs_begin > 0:
            raise AbrSdkUnexpectedState("Invalid replace_byte_offs_begin")
        if (
            c.replace_byte_offs_end > 0
            or c.replace_byte_offs_end < c.replace_byte_offs_begin
        ):
            raise AbrSdkUnexpectedState("Invalid replace_byte_offs_end")
        if c.n_bytes > len(c.data):
            raise AbrSdkUnexpectedState("Invalid n_bytes")
        return AsrChunk(
            type=AsrTextChunkType(c.type),
            replace_byte_offset_begin=c.replace_byte_offs_begin,
            replace_byte_offset_end=c.replace_byte_offs_end,
            data=bytes(c.data[: c.n_bytes]),
        )

    def update(self, buf: bytearray) -> None:
        """Apply this chunk to a running transcript ``bytearray``."""
        if self.type == AsrTextChunkType.NONE:
            return
        i0 = len(buf) + self.replace_byte_offset_begin
        i1 = len(buf) + self.replace_byte_offset_end
        if i0 < 0 or i1 < 0:
            raise ValueError("Buffer not compatible with this chunk")
        buf[:] = buf[:i0] + self.data + buf[i1:]


class AsrTranscript:
    """Collected ASR output chunks with text assembly."""

    def __init__(self) -> None:
        self.chunks: list[AsrChunk] = []

    @property
    def text(self) -> str:
        """Assemble and return the full transcript text from all chunks."""
        buf = bytearray()
        for chunk in self.chunks:
            chunk.update(buf)
        return buf.decode("utf-8")


class Processor:
    """
    Event loop for streaming PCM audio through an :class:`Asr` application.

    Feeds audio into the ASR pipeline and delivers transcribed text chunks.
    Attach to an :class:`Asr` instance and push audio data incrementally.
    Output chunks are delivered via the *on_chunk* callback::

        with Processor(asr, on_chunk=my_callback) as proc:
            proc.push(chunk1)
            proc.push(chunk2)
            proc.wait_for_completion()

    This class is also used internally by :meth:`Asr.process`.
    """

    def __init__(
        self,
        asr: Asr,
        on_chunk: Callable[[AsrChunk], None] | None = None,
    ) -> None:
        self._asr = asr
        self._on_chunk = on_chunk
        self._evset: EventSet | None = asr.create_event_set()

        # We want to be notified if there are at least two bytes of space (one
        # 16-bit sample) available in the input buffer. We generally want this
        # event to be disabled unless we're actively feeding data into the
        # buffer.
        self._input_event = self._evset.create_buffer_level_event(
            asr.input_buffer, 2, EventFlags(auto_disable=True, disabled=True)
        )

        # We want to be notified as soon as there is at least one AsrChunk
        # available on the output. This event shouldn't auto-disable; we're
        # always interested in output data becoming available.
        self._output_event = self._evset.create_buffer_level_event(
            asr.text_chunk_output_buffer, AsrChunk.SIZE
        )

        # We want to be notified as soon as the application becomes idle (all
        # NNs drained AND every processor's flush has reached its terminal
        # state); generally enabled manually when we're waiting for processing
        # to be done.
        self._idle_event = self._evset.create_application_idle_event(
            EventFlags(auto_disable=True, disabled=True)
        )

    def process_and_wait_for_output(
        self,
        data: bytes | None,
        timeout_ms: int,
        flush: bool,
    ) -> None:
        """
        Push input data and wait for output text chunks.

        This is the core event loop.  Higher-level methods :meth:`push` and
        :meth:`wait_for_completion` delegate to this method.

        Parameters
        ----------
        data
            PCM input bytes (little-endian 16-bit), or *None* to push no
            new data.
        timeout_ms
            Maximum time in milliseconds to spend waiting for output
            after all input has been pushed.  ``0`` means return
            immediately once input is consumed.  The timeout is measured
            from when this method is called.
        flush
            If *True*, flush the ASR pipeline after all input is streamed
            and wait until the neural network becomes idle.
        """
        if self._evset is None:
            raise RuntimeError("Processor is closed")

        # A memoryview makes the per-iteration buf[offset:] slices zero-copy.
        buf = memoryview(data if data is not None else b"")
        offset = 0
        did_flush = False
        t_start = time.monotonic_ns() // 1_000_000 if timeout_ms > 0 else 0

        while True:
            # Push as much input as possible.
            remaining = len(buf) - offset
            if remaining > 0:
                n_written = self._asr.input_buffer.push(buf[offset:])
                offset += n_written
                remaining = len(buf) - offset

            # Ensure that the "input buffer space event" is enabled.
            if remaining > 0:
                self._input_event.enable()

            # Flush once all input is consumed.
            if flush and remaining == 0 and not did_flush:
                self._asr.flush()
                self._idle_event.enable()
                did_flush = True

            # During the atomic phase (still feeding input or flushing),
            # use an infinite timeout.  Otherwise use the caller's value.
            adjusted_timeout = -1 if remaining > 0 or flush else timeout_ms

            # Account for time already spent in this call.
            if adjusted_timeout > 0:
                elapsed = time.monotonic_ns() // 1_000_000 - t_start
                adjusted_timeout = max(adjusted_timeout - elapsed, 0)

            # Poll for events.
            try:
                if not self._evset.poll(adjusted_timeout):
                    return  # Timeout expired without any event.
            except InterruptedError:
                # Only propagate the interrupt when we are past the
                # atomic input-feeding/flushing phase.
                if adjusted_timeout > 0 and timeout_ms > 0:
                    raise

            # Drain all complete chunks in the output buffer. The output
            # event is level-triggered and always enabled, so the next poll
            # wakes again if more chunks arrive.
            while self._output_event.is_triggered:
                chunk_bytes = self._asr.text_chunk_output_buffer.pull(AsrChunk.SIZE)
                if len(chunk_bytes) == 0:
                    break
                if len(chunk_bytes) != AsrChunk.SIZE:
                    raise AbrSdkUnexpectedState(
                        f"Pulled a partial text chunk ({len(chunk_bytes)} of "
                        f"{AsrChunk.SIZE} bytes) from the output buffer"
                    )
                chunk = AsrChunk.parse(chunk_bytes)
                if self._on_chunk is not None:
                    self._on_chunk(chunk)

            if flush and self._idle_event.is_triggered:
                flush = False

    def push(self, data: bytes, output_poll_timeout_ms: int = 0) -> None:
        """
        Push PCM audio data into the ASR network.

        This may block briefly if the input buffer is full.  The
        ``on_chunk`` callback may be invoked during this call.

        Parameters
        ----------
        data
            PCM audio as little-endian 16-bit bytes.
        output_poll_timeout_ms
            Extra time in milliseconds to spend waiting for output after
            the input has been consumed.  ``0`` (the default) returns as
            soon as the input is pushed.
        """
        if output_poll_timeout_ms < 0:
            raise ValueError("output_poll_timeout_ms must not be negative")
        self.process_and_wait_for_output(data, output_poll_timeout_ms, False)

    def wait_for_completion(self) -> None:
        """
        Block until all previously pushed data is fully processed.

        The ``on_chunk`` callback may be invoked during this call.
        """
        self.process_and_wait_for_output(None, 0, True)

    def close(self) -> None:
        """Release all event resources held by this processor."""
        if self._evset is None:
            return
        self._evset.close()
        self._evset = None

    def __enter__(self) -> Processor:
        return self

    def __exit__(
        self,
        type_: type[BaseException] | None,
        value: BaseException | None,
        traceback: TracebackType | None,
    ) -> None:
        self.close()
