#!/usr/bin/env python3
#
# Copyright (c) 2024 Stormux
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the
# Free Software Foundation, Inc., Franklin Street, Fifth Floor,
# Boston MA  02110-1301 USA.
#
# Cthulhu project: https://git.stormux.org/storm/cthulhu

"""Provides a Cthulhu speech server for Piper TTS backend."""

__id__ = "$Id$"
__version__ = "$Revision$"
__date__ = "$Date$"
__copyright__ = "Copyright (c) 2024 Stormux"
__license__ = "LGPL"

import io
import os
import sys
import threading
import time
import wave
from concurrent.futures import ThreadPoolExecutor

from gi.repository import GLib

from . import chnames
from . import debug
from . import guilabels
from . import messages
from . import settings
from . import speechserver
from . import cthulhu_state
from .acss import ACSS
from . import piper_voice_manager
from . import piper_audio_player

def _getPipxSitePackages():
    pipxHome = os.environ.get("PIPX_HOME", os.path.expanduser("~/.local/pipx"))
    pythonVersion = f"python{sys.version_info.major}.{sys.version_info.minor}"
    sitePackages = os.path.join(
        pipxHome,
        "venvs",
        "piper-tts",
        "lib",
        pythonVersion,
        "site-packages"
    )
    if os.path.isdir(sitePackages):
        return sitePackages
    return None


def _tryImportPiper():
    try:
        from piper.voice import PiperVoice
        return PiperVoice
    except ImportError:
        pass

    sitePackages = _getPipxSitePackages()
    if not sitePackages:
        return None

    addedPath = False
    if sitePackages not in sys.path:
        sys.path.insert(0, sitePackages)
        addedPath = True

    try:
        from piper.voice import PiperVoice
        msg = f'PIPER: Loaded piper-tts from pipx venv ({sitePackages})'
        debug.printMessage(debug.LEVEL_INFO, msg, True)
        return PiperVoice
    except ImportError:
        if addedPath:
            try:
                sys.path.remove(sitePackages)
            except ValueError:
                pass
        return None


PiperVoice = _tryImportPiper()
_piperAvailable = PiperVoice is not None
if _piperAvailable:
    try:
        from piper.config import SynthesisConfig as _PiperSynthesisConfig
    except Exception:
        _PiperSynthesisConfig = None
else:
    _PiperSynthesisConfig = None
if not _piperAvailable:
    msg = 'PIPER: piper-tts library not available'
    debug.printMessage(debug.LEVEL_INFO, msg, True)


class SpeechServer(speechserver.SpeechServer):
    """Piper TTS speech server implementation.

    Provides speech synthesis using Piper neural text-to-speech,
    implementing the Cthulhu SpeechServer interface.
    """

    _active_servers = {}

    DEFAULT_SERVER_ID = 'piper-default'

    @staticmethod
    def getFactoryName():
        """Returns a localized name describing this factory."""
        return guilabels.PIPER_TTS

    @staticmethod
    def getSpeechServers():
        """Gets available speech servers as a list."""
        servers = []

        if not _piperAvailable:
            msg = 'PIPER: Cannot list servers - piper-tts not installed'
            debug.printMessage(debug.LEVEL_INFO, msg, True)
            return servers

        manager = piper_voice_manager.getManager()
        voices = manager.discoverVoices()

        if not voices:
            msg = 'PIPER: No voice models found'
            debug.printMessage(debug.LEVEL_INFO, msg, True)
            return servers

        server = SpeechServer._getSpeechServer(SpeechServer.DEFAULT_SERVER_ID)
        if server is not None:
            servers.append(server)

        return servers

    @classmethod
    def _getSpeechServer(cls, serverId):
        """Return an active server for given id.

        Attempt to create the server if it doesn't exist yet.
        Returns None when it is not possible to create the server.
        """
        if serverId not in cls._active_servers:
            cls(serverId)
        return cls._active_servers.get(serverId)

    @staticmethod
    def getSpeechServer(info=None):
        """Gets a given SpeechServer based upon the info."""
        thisId = info[1] if info is not None else SpeechServer.DEFAULT_SERVER_ID
        return SpeechServer._getSpeechServer(thisId)

    @staticmethod
    def shutdownActiveServers():
        """Cleans up and shuts down this factory."""
        servers = list(SpeechServer._active_servers.values())
        for server in servers:
            server.shutdown()

    def __init__(self, serverId):
        """Initialize the Piper speech server.

        Arguments:
        - serverId: Identifier for this server instance
        """
        super(SpeechServer, self).__init__()
        self._id = serverId
        self._voice = None
        self._voiceInfo = None
        self._voiceManager = piper_voice_manager.getManager()
        self._audioPlayer = None
        self._executor = None
        self._currentFuture = None
        self._stopEvent = threading.Event()
        self._lock = threading.Lock()
        self._speakGeneration = 0

        self._currentVoiceProperties = {}
        self._acssManipulators = (
            (ACSS.RATE, self._setRate),
            (ACSS.AVERAGE_PITCH, self._setPitch),
            (ACSS.GAIN, self._setVolume),
            (ACSS.FAMILY, self._setFamily),
        )

        self._rate = 50
        self._pitch = 5.0
        self._volume = 1.0

        if not _piperAvailable:
            msg = 'PIPER: piper-tts library not available'
            debug.printMessage(debug.LEVEL_WARNING, msg, True)
            return

        try:
            self._init()
        except Exception as e:
            debug.printException(debug.LEVEL_WARNING)
            msg = f'PIPER: Failed to initialize server: {e}'
            debug.printMessage(debug.LEVEL_WARNING, msg, True)
        else:
            SpeechServer._active_servers[serverId] = self

    def _init(self):
        """Initialize the speech server components."""
        voices = self._voiceManager.getVoices()
        if not voices:
            raise RuntimeError("No Piper voices found")

        self._voiceInfo = voices[0]
        self._loadVoice(self._voiceInfo)

        self._audioPlayer = piper_audio_player.PiperAudioPlayer(
            self._voiceInfo.sampleRate
        )

        self._executor = ThreadPoolExecutor(
            max_workers=1,
            thread_name_prefix="piper"
        )

        msg = f'PIPER: Initialized with voice {self._voiceInfo.displayName}'
        debug.printMessage(debug.LEVEL_INFO, msg, True)

    def _loadVoice(self, voiceInfo):
        """Load a Piper voice model.

        Arguments:
        - voiceInfo: PiperVoiceInfo for the voice to load
        """
        msg = f'PIPER: Loading voice {voiceInfo.displayName}'
        debug.printMessage(debug.LEVEL_INFO, msg, True)

        self._voice = PiperVoice.load(str(voiceInfo.modelPath))
        self._voiceInfo = voiceInfo
        detectedRate = self._getVoiceSampleRate(self._voice, voiceInfo)
        if detectedRate and detectedRate != self._voiceInfo.sampleRate:
            msg = (
                f'PIPER: Using detected sample rate {detectedRate} '
                f'for {voiceInfo.displayName}'
            )
            debug.printMessage(debug.LEVEL_INFO, msg, True)
            self._voiceInfo.sampleRate = detectedRate

        if self._audioPlayer:
            self._audioPlayer.setSampleRate(voiceInfo.sampleRate)

    def _getVoiceSampleRate(self, voice, voiceInfo):
        if voice is None:
            return voiceInfo.sampleRate if voiceInfo else None

        for attr in ("sample_rate", "sampleRate"):
            value = getattr(voice, attr, None)
            if isinstance(value, (int, float)) and value > 0:
                return int(value)

        config = getattr(voice, "config", None)
        sampleRate = None
        if isinstance(config, dict):
            audio = config.get("audio")
            if isinstance(audio, dict):
                sampleRate = audio.get("sample_rate")
            elif audio is not None and hasattr(audio, "sample_rate"):
                sampleRate = getattr(audio, "sample_rate", None)
            if sampleRate is None:
                sampleRate = config.get("sample_rate")
        else:
            if hasattr(config, "audio"):
                audio = getattr(config, "audio")
                if isinstance(audio, dict):
                    sampleRate = audio.get("sample_rate")
                elif audio is not None and hasattr(audio, "sample_rate"):
                    sampleRate = getattr(audio, "sample_rate", None)
            if sampleRate is None and hasattr(config, "sample_rate"):
                sampleRate = getattr(config, "sample_rate", None)

        try:
            sampleRate = int(sampleRate) if sampleRate is not None else None
        except (TypeError, ValueError):
            sampleRate = None

        if sampleRate and sampleRate > 0:
            return sampleRate

        return voiceInfo.sampleRate if voiceInfo else None

    def _mapRate(self, acssRate):
        """Map ACSS rate (0-99) to Piper length_scale.

        ACSS rate 50 (default) = length_scale 1.0
        Higher ACSS rate = lower length_scale (faster)
        Lower ACSS rate = higher length_scale (slower)

        Arguments:
        - acssRate: Rate value from 0-99
        """
        rate = acssRate if acssRate is not None else 50
        rate = max(0, min(99, rate))
        lengthScale = 2.0 - (rate / 99.0) * 1.5
        return max(0.5, min(2.0, lengthScale))

    def _mapPitch(self, acssPitch):
        """Map ACSS pitch (0-9) to pitch adjustment factor.

        Note: Piper's native pitch control is limited.
        This maps to a factor that could be used for post-processing.

        Arguments:
        - acssPitch: Pitch value from 0-9
        """
        pitch = acssPitch if acssPitch is not None else 5.0
        pitch = max(0, min(9, pitch))
        return pitch

    def _mapVolume(self, acssGain):
        """Map ACSS gain (0-9) to volume (0.0-1.0).

        Arguments:
        - acssGain: Gain value from 0-9
        """
        gain = acssGain if acssGain is not None else 10
        gain = max(0, min(10, gain))
        return gain / 10.0

    def _setRate(self, acssRate):
        """Set the speech rate.

        Arguments:
        - acssRate: ACSS rate value (0-99)
        """
        self._rate = acssRate if acssRate is not None else 50
        msg = f'PIPER: Rate set to {self._rate}'
        debug.printMessage(debug.LEVEL_INFO, msg, True)

    def _setPitch(self, acssPitch):
        """Set the speech pitch.

        Arguments:
        - acssPitch: ACSS pitch value (0-9)
        """
        self._pitch = acssPitch if acssPitch is not None else 5.0
        msg = f'PIPER: Pitch set to {self._pitch}'
        debug.printMessage(debug.LEVEL_INFO, msg, True)

    def _setVolume(self, acssGain):
        """Set the speech volume.

        Arguments:
        - acssGain: ACSS gain value (0-9)
        """
        self._volume = self._mapVolume(acssGain)
        if self._audioPlayer:
            self._audioPlayer.setVolume(self._volume)
        msg = f'PIPER: Volume set to {self._volume}'
        debug.printMessage(debug.LEVEL_INFO, msg, True)

    def _setFamily(self, acssFamily):
        """Set the voice family.

        Arguments:
        - acssFamily: ACSS family dict with voice info
        """
        if not acssFamily:
            return

        name = acssFamily.get(speechserver.VoiceFamily.NAME)
        if not name:
            return

        for voice in self._voiceManager.getVoices():
            if voice.displayName == name:
                if voice != self._voiceInfo:
                    try:
                        self._loadVoice(voice)
                    except Exception as e:
                        msg = f'PIPER: Failed to load voice {name}: {e}'
                        debug.printMessage(debug.LEVEL_WARNING, msg, True)
                break

    def _applyAcss(self, acss):
        """Apply ACSS voice settings.

        Arguments:
        - acss: ACSS settings to apply
        """
        if acss is None:
            acss = settings.voices[settings.DEFAULT_VOICE]

        with self._lock:
            current = self._currentVoiceProperties
            for acssProperty, method in self._acssManipulators:
                value = acss.get(acssProperty)
                if value is not None:
                    if current.get(acssProperty) != value:
                        method(value)
                        current[acssProperty] = value
                elif acssProperty == ACSS.AVERAGE_PITCH:
                    method(5.0)
                    current[acssProperty] = 5.0
                elif acssProperty == ACSS.GAIN:
                    method(10)
                    current[acssProperty] = 10
                elif acssProperty == ACSS.RATE:
                    method(50)
                    current[acssProperty] = 50

    def _synthesize(self, text):
        """Synthesize text to audio data.

        Arguments:
        - text: Text to synthesize

        Returns raw PCM audio data as bytes.
        """
        if not self._voice or not text:
            return None

        lengthScale = self._mapRate(self._rate)

        if _PiperSynthesisConfig is not None:
            try:
                synConfig = _PiperSynthesisConfig(length_scale=lengthScale)
                audioChunks = self._voice.synthesize(text, syn_config=synConfig)
                audioParts = []
                sampleRate = None
                for chunk in audioChunks:
                    if sampleRate is None:
                        sampleRate = chunk.sample_rate
                    audioParts.append(chunk.audio_int16_bytes)
                if sampleRate and sampleRate != self._voiceInfo.sampleRate:
                    self._voiceInfo.sampleRate = sampleRate
                    if self._audioPlayer:
                        self._audioPlayer.setSampleRate(sampleRate)
                audioData = b"".join(audioParts)
                return audioData if audioData else None
            except TypeError:
                pass

        wavBuffer = io.BytesIO()
        with wave.open(wavBuffer, 'wb') as wavFile:
            wavFile.setnchannels(1)
            wavFile.setsampwidth(2)
            wavFile.setframerate(self._voiceInfo.sampleRate)

            self._voice.synthesize(
                text,
                wavFile,
                length_scale=lengthScale
            )

        wavBuffer.seek(44)
        return wavBuffer.read()

    def _synthesizeAndPlay(self, text, acss, completionCallback=None, generation=0):
        """Synthesize and play text (runs in worker thread).

        Arguments:
        - text: Text to synthesize
        - acss: ACSS settings
        - completionCallback: Optional callback when complete
        """
        try:
            if self._stopEvent.is_set() or generation != self._speakGeneration:
                return

            self._applyAcss(acss)

            audioData = self._synthesize(text)
            if not audioData or self._stopEvent.is_set() or generation != self._speakGeneration:
                return

            if self._audioPlayer:
                self._audioPlayer.play(audioData, False)
                while self._audioPlayer.isPlaying():
                    if self._stopEvent.is_set() or generation != self._speakGeneration:
                        self._audioPlayer.stop()
                        return
                    time.sleep(0.01)

            if completionCallback and not self._stopEvent.is_set() and generation == self._speakGeneration:
                GLib.idle_add(completionCallback)

        except Exception as e:
            msg = f'PIPER: Synthesis error: {e}'
            debug.printMessage(debug.LEVEL_WARNING, msg, True)

    def getInfo(self):
        """Returns [name, id]."""
        return [guilabels.PIPER_TTS, self._id]

    def getVoiceFamilies(self):
        """Returns a list of VoiceFamily instances."""
        families = []
        for voice in self._voiceManager.getVoices():
            family = self._voiceManager.voiceToVoiceFamily(voice)
            families.append(family)
        return families

    def speak(self, text=None, acss=None, interrupt=True):
        """Speaks the given text.

        Arguments:
        - text: Text to speak
        - acss: ACSS voice settings
        - interrupt: If True, stop any current speech first
        """
        if not text or not self._voice:
            return

        if interrupt:
            self.stop()

        with self._lock:
            self._stopEvent.clear()
            generation = self._speakGeneration

        msg = f"PIPER: Speaking '{text}'"
        debug.printMessage(debug.LEVEL_INFO, msg, True)

        self._currentFuture = self._executor.submit(
            self._synthesizeAndPlay, text, acss, None, generation
        )

    def sayAll(self, utteranceIterator, progressCallback):
        """Iterates through utterances, speaking each one.

        Arguments:
        - utteranceIterator: Iterator yielding [SayAllContext, acss] tuples
        - progressCallback: Called with progress updates
        """
        GLib.idle_add(self._sayAllWorker, utteranceIterator, progressCallback)

    def _sayAllWorker(self, iterator, callback):
        """Process one utterance at a time (called via GLib.idle_add).

        Arguments:
        - iterator: Utterance iterator
        - callback: Progress callback
        """
        try:
            context, acss = next(iterator)
        except StopIteration:
            return False

        def onComplete():
            context.currentOffset = context.endOffset
            callback(context.copy(), speechserver.SayAllContext.COMPLETED)
            GLib.idle_add(self._sayAllWorker, iterator, callback)

        context.currentOffset = context.startOffset
        callback(context.copy(), speechserver.SayAllContext.PROGRESS)

        with self._lock:
            self._stopEvent.clear()
            generation = self._speakGeneration
        self._currentFuture = self._executor.submit(
            self._synthesizeAndPlay, context.utterance, acss, onComplete, generation
        )

        return False

    def speakCharacter(self, character, acss=None):
        """Speaks a single character immediately.

        Arguments:
        - character: Character to speak
        - acss: ACSS voice settings
        """
        name = chnames.getCharacterName(character)
        if name and name != character:
            if cthulhu_state.activeScript:
                name = cthulhu_state.activeScript.utilities.adjustForPronunciation(name)
            self.speak(name, acss)
        else:
            self.speak(character, acss)

    def speakKeyEvent(self, event, acss=None):
        """Speaks a key event immediately.

        Arguments:
        - event: The KeyboardEvent to speak
        - acss: ACSS voice settings
        """
        eventString = event.getKeyName()
        lockingStateString = event.getLockingStateString()
        eventString = f"{eventString} {lockingStateString}".strip()
        self.speak(eventString, acss)

    def _changeDefaultSpeechRate(self, step, decrease=False):
        """Change the default speech rate.

        Arguments:
        - step: Amount to change
        - decrease: If True, decrease rate; otherwise increase
        """
        acss = settings.voices[settings.DEFAULT_VOICE]
        delta = step * (-1 if decrease else 1)
        try:
            rate = acss[ACSS.RATE]
        except KeyError:
            rate = 50
        acss[ACSS.RATE] = max(0, min(99, rate + delta))
        msg = f"PIPER: Rate set to {acss[ACSS.RATE]}"
        debug.printMessage(debug.LEVEL_INFO, msg, True)
        self.speak(
            messages.SPEECH_SLOWER if decrease else messages.SPEECH_FASTER,
            acss=acss
        )

    def _changeDefaultSpeechPitch(self, step, decrease=False):
        """Change the default speech pitch.

        Arguments:
        - step: Amount to change
        - decrease: If True, decrease pitch; otherwise increase
        """
        acss = settings.voices[settings.DEFAULT_VOICE]
        delta = step * (-1 if decrease else 1)
        try:
            pitch = acss[ACSS.AVERAGE_PITCH]
        except KeyError:
            pitch = 5
        acss[ACSS.AVERAGE_PITCH] = max(0, min(9, pitch + delta))
        msg = f"PIPER: Pitch set to {acss[ACSS.AVERAGE_PITCH]}"
        debug.printMessage(debug.LEVEL_INFO, msg, True)
        self.speak(
            messages.SPEECH_LOWER if decrease else messages.SPEECH_HIGHER,
            acss=acss
        )

    def _changeDefaultSpeechVolume(self, step, decrease=False):
        """Change the default speech volume.

        Arguments:
        - step: Amount to change
        - decrease: If True, decrease volume; otherwise increase
        """
        acss = settings.voices[settings.DEFAULT_VOICE]
        delta = step * (-1 if decrease else 1)
        try:
            volume = acss[ACSS.GAIN]
        except KeyError:
            volume = 10
        acss[ACSS.GAIN] = max(0, min(9, volume + delta))
        msg = f"PIPER: Volume set to {acss[ACSS.GAIN]}"
        debug.printMessage(debug.LEVEL_INFO, msg, True)
        self.speak(
            messages.SPEECH_SOFTER if decrease else messages.SPEECH_LOUDER,
            acss=acss
        )

    def increaseSpeechRate(self, step=5):
        """Increases the speech rate."""
        self._changeDefaultSpeechRate(step)

    def decreaseSpeechRate(self, step=5):
        """Decreases the speech rate."""
        self._changeDefaultSpeechRate(step, decrease=True)

    def increaseSpeechPitch(self, step=0.5):
        """Increases the speech pitch."""
        self._changeDefaultSpeechPitch(step)

    def decreaseSpeechPitch(self, step=0.5):
        """Decreases the speech pitch."""
        self._changeDefaultSpeechPitch(step, decrease=True)

    def increaseSpeechVolume(self, step=0.5):
        """Increases the speech volume."""
        self._changeDefaultSpeechVolume(step)

    def decreaseSpeechVolume(self, step=0.5):
        """Decreases the speech volume."""
        self._changeDefaultSpeechVolume(step, decrease=True)

    def updateCapitalizationStyle(self):
        """Updates the capitalization style used by the speech server."""
        pass

    def updatePunctuationLevel(self):
        """Punctuation level changed, inform this speechServer."""
        pass

    def stop(self):
        """Stops ongoing speech and flushes the queue."""
        with self._lock:
            self._speakGeneration += 1
            self._stopEvent.set()

        if self._currentFuture:
            self._currentFuture.cancel()
            self._currentFuture = None

        if self._audioPlayer:
            self._audioPlayer.stop()

    def shutdown(self):
        """Shuts down the speech engine."""
        self.stop()

        if self._executor:
            self._executor.shutdown(wait=False)
            self._executor = None

        if self._audioPlayer:
            self._audioPlayer.shutdown()
            self._audioPlayer = None

        self._voice = None
        self._voiceInfo = None

        if self._id in SpeechServer._active_servers:
            del SpeechServer._active_servers[self._id]

        msg = 'PIPER: Server shutdown complete'
        debug.printMessage(debug.LEVEL_INFO, msg, True)

    def reset(self, text=None, acss=None):
        """Resets the speech engine."""
        self.stop()
        if self._voiceInfo:
            try:
                self._loadVoice(self._voiceInfo)
            except Exception as e:
                msg = f'PIPER: Failed to reset voice: {e}'
                debug.printMessage(debug.LEVEL_WARNING, msg, True)
