# (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.

"""
API for running TTS (text-to-speech) applications.

Author: Pawel Jaworski
"""

from __future__ import annotations

import contextlib
import logging
import sys
import time
from collections.abc import Callable
from pathlib import Path
from typing import Any

from abr_sdk.core import Application, Buffer, EventFlags, Library
from abr_sdk.exceptions import AbrSdkError
from abr_sdk.tts_preprocess import find_unpreprocessed_chars


class Tts(Application):
    """
    Synthesize speech audio from input text.

    Push raw UTF-8 text bytes and receive synthesized mono PCM audio,
    one signed 16-bit integer per sample in little-endian order at
    16 kHz.

    Streaming usage::

        chunks: list[bytes] = []
        with Tts("libabr-tts.so") as tts:
            tts.push(b"Hello world.", on_pcm=chunks.append)
            tts.wait_for_completion()
        pcm = b"".join(chunks)
    """

    def __init__(
        self,
        lib_or_path: str | Path | Library,
        *,
        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:
        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="tts",
            resources_dir=resources_dir,
            logger=logger,
        )

        # Free the native handle if the TTS-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_tts_available(self._handle):
                raise AbrSdkError("The loaded application does not support TTS")
            self.text_input_buffer = Buffer(
                self._cabi, self._cabi.abr_tts_get_text_input_buffer(self._handle)
            )
            self.pcm_output_buffer = Buffer(
                self._cabi, self._cabi.abr_tts_get_pcm_output_buffer(self._handle)
            )
            close_on_error.pop_all()
        self._processor: Processor | None = None
        self._preprocess_warned = False

    def __enter__(self) -> Tts:
        return self

    def push(
        self,
        data: bytes,
        *,
        on_pcm: Callable[[bytes], None] | None = None,
        output_poll_timeout_ms: int = 0,
    ) -> None:
        """
        Push UTF-8 text bytes into the TTS pipeline.

        Advances the event loop until all input has been accepted by the
        pipeline. PCM bytes produced while advancing are delivered to
        ``on_pcm``.
        """
        self._ensure_alive()
        if output_poll_timeout_ms < 0:
            raise ValueError("output_poll_timeout_ms must not be negative")
        self._warn_if_not_preprocessed(data)
        if self._processor is None:
            self._processor = Processor(self, on_pcm=on_pcm)
        self._processor.process_and_wait_for_output(data, output_poll_timeout_ms, False)

    def _warn_if_not_preprocessed(self, data: bytes) -> None:
        """
        Warn once if pushed text was not run through ``tts_preprocess``.

        Raw digits or symbols reach the phoneme dictionary and get mangled or
        dropped. Best-effort: one warning per instance, on stderr so it never
        corrupts PCM written to stdout.
        """
        if self._preprocess_warned:
            return
        offenders = find_unpreprocessed_chars(data.decode("utf-8", errors="replace"))
        if len(offenders) == 0:
            return
        self._preprocess_warned = True
        shown = ", ".join(repr(c) for c in sorted(offenders))
        print(
            f"TTS warning: pushed text contains characters that "
            f"tts_preprocess() removes ({shown}); call "
            f"abr_sdk.tts_preprocess.tts_preprocess() on the text before "
            f"push(), or the backend will mispronounce or drop it.",
            file=sys.stderr,
        )

    def wait_for_completion(self) -> None:
        """
        Block until the TTS pipeline finishes synthesizing all queued text.

        Marks the end of input so the pipeline flushes its networks, then
        drives the event loop until the application reports idle (every NN
        has consumed its input and produced nothing more).
        """
        self._ensure_alive()
        if self._processor is None:
            return
        try:
            self._processor.process_and_wait_for_output(b"\x00", 0, True)
        finally:
            self._processor.close()
            self._processor = None

    def close(self) -> None:
        """Close the TTS instance and release any active processor."""
        if self._closed:
            return
        proc = getattr(self, "_processor", None)
        if proc is not None:
            proc.close()
            self._processor = None
        super().close()


class Processor:
    """Event loop driving a streaming TTS session."""

    def __init__(
        self,
        tts: Tts,
        on_pcm: Callable[[bytes], None] | None = None,
    ) -> None:
        self._tts = tts
        self._on_pcm = on_pcm
        self._evset = tts.create_event_set()
        self._input_event = self._evset.create_buffer_level_event(
            tts.text_input_buffer, 1, EventFlags(auto_disable=True)
        )
        # Wake as soon as any PCM is available (one 16-bit sample).
        self._output_event = self._evset.create_buffer_level_event(
            tts.pcm_output_buffer, 2
        )
        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,
        drain_to_idle: bool,
    ) -> None:
        """
        Push ``data`` into TTS and advance the event loop until input is drained.

        When ``drain_to_idle`` is true, keeps polling past input
        exhaustion until the application-idle event fires (all NNs
        finished consuming and produced no further PCM).
        """
        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
        idle_enabled = 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._tts.text_input_buffer.push(buf[offset:])
                offset += n_written
                remaining = len(buf) - offset

            if remaining > 0:
                self._input_event.enable()

            # Wait for idle once all input is consumed.
            if drain_to_idle and remaining == 0 and not idle_enabled:
                self._idle_event.enable()
                idle_enabled = True

            adjusted_timeout = -1 if remaining > 0 or drain_to_idle else timeout_ms
            if adjusted_timeout > 0:
                elapsed = time.monotonic_ns() // 1_000_000 - t_start
                adjusted_timeout = max(adjusted_timeout - elapsed, 0)

            try:
                if not self._evset.poll(adjusted_timeout):
                    return
            except InterruptedError:
                if adjusted_timeout > 0 and timeout_ms > 0:
                    raise

            # Drain all available PCM. The output event is level-triggered
            # and always enabled, so the next poll wakes again if more
            # arrives.
            while self._output_event.is_triggered:
                pcm_bytes = self._tts.pcm_output_buffer.pull(
                    self._tts.pcm_output_buffer.size
                )
                if not pcm_bytes:
                    break
                if self._on_pcm is not None:
                    self._on_pcm(pcm_bytes)

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

    def close(self) -> None:
        """Release the underlying event set; safe to call multiple times."""
        if self._evset is None:
            return
        self._evset.close()
        self._evset = None

    def __enter__(self) -> Processor:
        return self

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