#!/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

"""Voice discovery and management for Piper TTS."""

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

import json
import os
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional

from . import debug
from . import speechserver


@dataclass
class PiperVoiceInfo:
    """Metadata for a Piper voice model."""

    name: str
    language: str
    dialect: str
    quality: str
    modelPath: Path
    configPath: Path
    sampleRate: int = 22050
    speakers: Dict[int, str] = field(default_factory=dict)

    @property
    def key(self) -> str:
        """Return a unique key for this voice."""
        dialectPart = f"-{self.dialect}" if self.dialect else ""
        return f"{self.language}{dialectPart}-{self.name}-{self.quality}"

    @property
    def displayName(self) -> str:
        """Return a human-readable display name."""
        dialectPart = f" ({self.dialect})" if self.dialect else ""
        return f"{self.name} - {self.language}{dialectPart} [{self.quality}]"

    @property
    def isMultiSpeaker(self) -> bool:
        """Return True if this is a multi-speaker model."""
        return len(self.speakers) > 1


class PiperVoiceManager:
    """Discovers and manages Piper voice models.

    Searches standard paths for Piper voice models (.onnx files with
    companion .onnx.json config files) and provides methods to list
    and load them.
    """

    VOICE_SEARCH_PATHS = [
        "~/.local/share/piper/voices",
        "~/.local/share/piper-voices",
        "~/.local/share/piper-tts/voices",
        "~/.config/piper/voices",
        "~/.config/piper-tts/voices",
        "$XDG_DATA_HOME/piper/voices",
        "$XDG_DATA_HOME/piper-voices",
        "$XDG_DATA_HOME/piper-tts/voices",
        "$XDG_DATA_HOME/cthulhu/piper-voices",
        "/usr/share/piper/voices",
        "/usr/share/piper-voices",
        "/usr/share/piper-tts/voices",
        "/usr/local/share/piper/voices",
        "/usr/local/share/piper-voices",
        "/usr/local/share/piper-tts/voices",
    ]

    VOICE_FILENAME_PATTERN = re.compile(
        r'^(?P<lang>[a-z]{2})(?:_(?P<dialect>[A-Z]{2}))?'
        r'-(?P<name>[a-zA-Z0-9_]+)'
        r'-(?P<quality>low|medium|high|x_low)\.onnx$'
    )

    def __init__(self, customPath=None):
        """Initialize the voice manager.

        Arguments:
        - customPath: Optional additional path to search for voices
        """
        self._customPath = customPath
        self._voices = []
        self._voiceCache = {}

    def discoverVoices(self) -> List[PiperVoiceInfo]:
        """Discover all available Piper voices.

        Searches standard paths and returns a list of PiperVoiceInfo
        objects for each valid voice found.
        """
        self._voices = []
        searchPaths = self._getSearchPaths()

        for searchPath in searchPaths:
            path = Path(os.path.expandvars(os.path.expanduser(searchPath)))
            if not path.exists():
                continue

            msg = f'PIPER VOICES: Searching {path}'
            debug.printMessage(debug.LEVEL_INFO, msg, True)

            for onnxFile in path.rglob("*.onnx"):
                configFile = Path(str(onnxFile) + ".json")
                if not configFile.exists():
                    configFile = onnxFile.with_suffix(".onnx.json")

                if configFile.exists():
                    try:
                        voice = self._parseVoice(onnxFile, configFile)
                        if voice:
                            self._voices.append(voice)
                            msg = f'PIPER VOICES: Found voice {voice.displayName}'
                            debug.printMessage(debug.LEVEL_INFO, msg, True)
                    except Exception as e:
                        msg = f'PIPER VOICES: Failed to parse {onnxFile}: {e}'
                        debug.printMessage(debug.LEVEL_WARNING, msg, True)
                else:
                    voice = self._parseVoiceFromFilename(onnxFile)
                    if voice:
                        self._voices.append(voice)
                        msg = f'PIPER VOICES: Found voice {voice.displayName} (no config)'
                        debug.printMessage(debug.LEVEL_INFO, msg, True)

        self._voices.sort(key=lambda v: (v.language, v.name, v.quality))

        msg = f'PIPER VOICES: Discovered {len(self._voices)} voice(s)'
        debug.printMessage(debug.LEVEL_INFO, msg, True)

        return self._voices

    def getVoices(self) -> List[PiperVoiceInfo]:
        """Get the list of discovered voices.

        Returns cached list; call discoverVoices() first to refresh.
        """
        if not self._voices:
            self.discoverVoices()
        return self._voices

    def getVoiceByKey(self, key: str) -> Optional[PiperVoiceInfo]:
        """Get a voice by its unique key.

        Arguments:
        - key: Voice key (e.g., "en_US-lessac-medium")
        """
        for voice in self.getVoices():
            if voice.key == key:
                return voice
        return None

    def getVoiceByName(self, name: str) -> Optional[PiperVoiceInfo]:
        """Get a voice by name (first match).

        Arguments:
        - name: Voice name (e.g., "lessac")
        """
        for voice in self.getVoices():
            if voice.name == name:
                return voice
        return None

    def getVoicesForLanguage(self, lang: str) -> List[PiperVoiceInfo]:
        """Get all voices for a specific language.

        Arguments:
        - lang: Language code (e.g., "en")
        """
        return [v for v in self.getVoices() if v.language == lang]

    def voiceToVoiceFamily(self, voiceInfo: PiperVoiceInfo) -> speechserver.VoiceFamily:
        """Convert a PiperVoiceInfo to a Cthulhu VoiceFamily.

        Arguments:
        - voiceInfo: The PiperVoiceInfo to convert
        """
        return speechserver.VoiceFamily({
            speechserver.VoiceFamily.NAME: voiceInfo.displayName,
            speechserver.VoiceFamily.LANG: voiceInfo.language,
            speechserver.VoiceFamily.DIALECT: voiceInfo.dialect or "",
            speechserver.VoiceFamily.VARIANT: voiceInfo.quality,
        })

    def _getSearchPaths(self) -> List[str]:
        """Get the list of paths to search for voices."""
        paths = list(self.VOICE_SEARCH_PATHS)

        if self._customPath:
            paths.insert(0, self._customPath)

        xdgDataHome = os.environ.get("XDG_DATA_HOME", "~/.local/share")
        paths = [p.replace("$XDG_DATA_HOME", xdgDataHome) for p in paths]

        return paths

    def _parseVoice(self, modelPath: Path, configPath: Path) -> Optional[PiperVoiceInfo]:
        """Parse a voice from its model and config files.

        Arguments:
        - modelPath: Path to the .onnx model file
        - configPath: Path to the .json config file
        """
        try:
            with open(configPath, 'r', encoding='utf-8') as f:
                config = json.load(f)
        except (json.JSONDecodeError, IOError) as e:
            msg = f'PIPER VOICES: Failed to read config {configPath}: {e}'
            debug.printMessage(debug.LEVEL_WARNING, msg, True)
            return None

        language = ""
        dialect = ""
        name = modelPath.stem
        quality = "medium"
        sampleRate = 22050
        speakers = {}

        if "language" in config:
            langInfo = config["language"]
            if isinstance(langInfo, dict):
                language = langInfo.get("code", "")
                if "_" in language:
                    parts = language.split("_")
                    language = parts[0]
                    dialect = parts[1] if len(parts) > 1 else ""
            elif isinstance(langInfo, str):
                if "_" in langInfo:
                    parts = langInfo.split("_")
                    language = parts[0]
                    dialect = parts[1] if len(parts) > 1 else ""
                else:
                    language = langInfo

        if "audio" in config:
            sampleRate = config["audio"].get("sample_rate", 22050)

        if "speaker_id_map" in config:
            speakers = {v: k for k, v in config["speaker_id_map"].items()}

        match = self.VOICE_FILENAME_PATTERN.match(modelPath.name)
        if match:
            if not language:
                language = match.group("lang")
            if not dialect:
                dialect = match.group("dialect") or ""
            name = match.group("name")
            quality = match.group("quality").replace("x_low", "x-low")
        else:
            stem = modelPath.stem
            for q in ["low", "medium", "high", "x_low", "x-low"]:
                if stem.endswith(f"-{q}"):
                    quality = q.replace("x_low", "x-low")
                    stem = stem[:-len(q)-1]
                    break

            parts = stem.split("-")
            if parts:
                langPart = parts[0]
                if "_" in langPart:
                    langParts = langPart.split("_")
                    if not language:
                        language = langParts[0]
                    if not dialect:
                        dialect = langParts[1] if len(langParts) > 1 else ""
                elif not language:
                    language = langPart

                if len(parts) > 1:
                    name = parts[1]

        if not language:
            language = "unknown"

        return PiperVoiceInfo(
            name=name,
            language=language,
            dialect=dialect,
            quality=quality,
            modelPath=modelPath,
            configPath=configPath,
            sampleRate=sampleRate,
            speakers=speakers
        )

    def _parseVoiceFromFilename(self, modelPath: Path) -> Optional[PiperVoiceInfo]:
        """Parse voice info from filename only (no config file).

        Arguments:
        - modelPath: Path to the .onnx model file
        """
        match = self.VOICE_FILENAME_PATTERN.match(modelPath.name)
        if not match:
            return None

        return PiperVoiceInfo(
            name=match.group("name"),
            language=match.group("lang"),
            dialect=match.group("dialect") or "",
            quality=match.group("quality").replace("x_low", "x-low"),
            modelPath=modelPath,
            configPath=modelPath.with_suffix(".onnx.json"),
            sampleRate=22050,
            speakers={}
        )


_manager = None


def getManager(customPath=None) -> PiperVoiceManager:
    """Get the singleton voice manager instance.

    Arguments:
    - customPath: Optional custom path to search for voices
    """
    global _manager
    if _manager is None:
        _manager = PiperVoiceManager(customPath)
    return _manager
