import argparse
import json
import logging
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional

from pydub import AudioSegment
from tinydb import TinyDB

log = logging.getLogger(__name__)


@dataclass
class Sentence:
    sentence: str
    sayable: List[str]
    unsayable: List[str]
    audio: Optional[AudioSegment] = None


class Voice:
    """Comprises all information and methods
    needed to index a folder of voice audio files
    and generate audio from them given a sentence string.
    """

    def __init__(self, path: Path, db_path: Path):
        """
        Args:
            path (Path): Path to folder of voice audio files.
            db_path (Path): Path to find or create database at.
        """
        self.path = path

        self.info_path = self.path.joinpath("info/")
        self.info_name = "info.json"

        self.db_path = Path(db_path)
        self.db_name = "db.json"

        self.name = self.path.name

        self._punctuation = [",", "."]
        self._punctuation_timing = {",": 250, ".": 500}

        self.alarms: List[str] = []

        self._db = self._setup_db(self.db_path, self.db_name)

        self._word_dict = self._build_word_dict(self.path)
        self._audio_format = self._find_audio_format(
            self._word_dict)  # TODO: Use properies?

        self._read_info(self.info_path, self.info_name)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self._db.close()

    # TODO: Need a better way to handle this
    def exit(self):
        self._db.close()

    def _setup_db(self, path: Path, db_name: str) -> TinyDB:
        """Sets up or finds existing database.

        Args:
            path (Path): Path to find or create database at. 
            db_name (str): Name of database file.

        Returns:
            TinyDB: Database.
        """
        path.mkdir(parents=True, exist_ok=True)
        return TinyDB(path.joinpath(db_name))

    def _build_word_dict(self, path: Path) -> Dict[str, Path]:
        """Builds dictionary of all available words.

        Args:
            path (Path): Path to folder of voice audio files.

        Raises:
            ValueError: Raised if there are duplicate filenames present.
            Exception: Raised if no words are found.

        Returns:
            Dict[str, Path]: Dict of {filepath: word} associations.
        """
        word_dict = {}
        word_files = list(x for x in path.iterdir() if x.is_file())

        for word in word_files:
            word = Path(word)
            name = str(word.stem).lower()
            if name in word_dict:
                raise ValueError(f"Word {name} is duplicated")
            word_dict[name] = word

        if len(word_dict) == 0:
            log.error("No words found")
            # TODO: Better default exception, or make some custom ones
            raise Exception("No words found")

        return word_dict

    def _read_info(self, path: Path, info_name: str):
        """Reads info file (if it exists)
        which may contain information about special
        groupings of words (like alarms)

        Args:
            path (Path): Path where info file resides.
            info_name (str): Name of info file.
        """
        # TODO: Allow arbitrary groupings of words
        info_path = path.joinpath(info_name)
        if info_path.exists():
            with open(info_path, 'r') as info_file:
                info = json.load(info_file)
                self.alarms = info["alarms"]
        log.info(f"Available alarms are {self.alarms}")

    def _find_audio_format(self, word_dict: Dict[str, Path]) -> str:
        """Determines audio format of voice audio files.

        Args:
            word_dict (Dict[str, Path]): Dict of {filepath: word} associations.

        Raises:
            Exception: Raised if no audio format can be determined,
                or if there are inconsistent audio formats.

        Returns:
            str: Audio format.
        """
        file_format = None
        for path in word_dict.values():
            if file_format == None:
                file_format = path.suffix[1:]
            else:
                if str(file_format) != str(path.suffix[1:]):
                    log.error(
                        "Inconsistent audio formats in the word dict")
                    # TODO: Better default exception, or make some custom ones
                    raise Exception("Inconsistent Audio Formats")
        if not file_format:
            # TODO: Better default exception, or make some custom ones
            raise Exception("No file format found")
        log.info(f"Audio format found: {file_format}")
        return file_format

    def _process_sentence(self, sentence: str) -> List[str]:
        """
        Takes a normally formatted sentence and breaks it into base elements

        Args:
            sentence (string): normally formatted sentence to be processed

        Returns:
            List[str]: array of elements in sentence
        """
        # TODO: This could use some rethinking. Should be easier to just always break punctuation marks
        # into their own elements, rather than selectively only dealing with trailing ones.
        split_sent = sentence.lower().rstrip().split(" ")
        log.info(f"Processing sentence '{sentence}'")

        # Pull out punctuation
        reduced_sent = []
        for item in split_sent:
            # find first punctuation mark, if any
            first_punct: Optional[str] = None
            try:
                first_punct = next(
                    (punct for punct in self._punctuation if punct in item))
            except StopIteration:
                pass

            if first_punct:
                # Get its index
                first_punct_ind = item.find(first_punct)

                # Add everything before punct (the word)
                reduced_sent.append(item[:first_punct_ind])

                # Add all the punctuation if its actually punctuation
                # TODO: Figure out if I want to deal with types like ".hello" throwing out all the characters after the period.
                for punct in item[first_punct_ind:]:
                    if punct in self._punctuation:
                        reduced_sent.append(punct)

            else:
                reduced_sent.append(item)

        # Clean blanks from reduced_sent
        if '' in reduced_sent:
            reduced_sent = [value for value in reduced_sent if value != '']

        log.info(f"Sentence processed: '{reduced_sent}'")
        return reduced_sent

    def _get_sayable_words(self, processed_sentence_arr: List[str], word_dict: Dict[str, Path]) -> List[str]:
        """Gets words in a sentence array that are available to be "said."

        Args:
            processed_sentence_arr (List[str]): Array of words in requested sentence.
            word_dict (Dict[str, Path]): Dict of {filepath: word} associations.

        Returns:
            List[str]: Sayable words.
        """
        sayable = [
            word for word in processed_sentence_arr if word in word_dict or word in self._punctuation]
        sayable_dedup = list(dict.fromkeys(sayable))  # removes duplicates
        sayable_dedup.sort()
        return sayable_dedup

    def _get_unsayable_words(self, processed_sentence_arr: List[str], sayable_words: List[str]) -> List[str]:
        """Gets words in a sentence array that are NOT available to be "said."

        Args:
            processed_sentence_arr (List[str]): Array of words in requested sentence.
            word_dict (List[str]): Dict of {filepath: word} associations.

        Returns:
            List[str]: Unsayable words.
        """
        unsayable = list(set(processed_sentence_arr) - set(sayable_words))
        return list(dict.fromkeys(unsayable))  # removes duplicates

    def _get_sayable_sentence_arr(self, processed_sentence_arr: List[str], sayable_words: List[str]) -> List[str]:
        """Removes words from sentence array that are not sayable.

        Args:
            processed_sentence_arr (List[str]): Array of words in sentence, in order.
            sayable_words (List[str]): Words from sentence that can actually be said.

        Returns:
            List[str]: Words in sentence that are sayable, in order.
        """
        # TODO: This is just a simple set operation. Function probably isn't needed. At least change to using a set.
        return [word for word in processed_sentence_arr if word in sayable_words]

    def _create_sentence_string(self, words: List[str]) -> str:
        """Joins sentence array into a string.

        Args:
            words (List[str]): Words in sentence, in order.

        Returns:
            str: Sentence string.
        """
        if len(words) == 1:
            return words[0]
        else:
            return " ".join(words)

    def get_words(self) -> List[str]:
        """Gets the available words for the voice

        Returns:
            List[str]: Words available to the voice
        """
        word_list = list(self._word_dict.keys())
        word_list.sort()
        return word_list

    def get_audio_format(self) -> str:
        """Get the audio format of the voice files as well as generated files
        Returns:
            (string): Audio format
        """
        return self._audio_format

    def get_sayable(self, sentence: str) -> List[str]:
        """Gets the words that the voice can say from a sentence

        Args:
            sentence (string): Sentence to check for sayable words
        Returns:
            List[str]: Words that can be said
        """

        proc_sentence = self._process_sentence(sentence)
        return self._get_sayable_words(proc_sentence, self._word_dict)

    def get_unsayable(self, sentence: str) -> List[str]:
        """Gets the words that the voice cannot say from a sentence

        Args:
            sentence (string): Sentence to check for un-sayable words
        Returns:
            List[str]: Words that cannot be said
        """

        proc_sentence = self._process_sentence(sentence)
        sayable_words = self._get_sayable_words(proc_sentence, self._word_dict)
        return self._get_unsayable_words(proc_sentence, sayable_words)

    def get_sentence_list(self, sentence: str) -> List[str]:
        """Converts sentence into list of words that can be said

        Args:
            sentence (string): Sentence to convert to list
        Returns:
            List[str]: Sentence in list format excluding words that cannot be said
        """

        proc_sentence = self._process_sentence(sentence)
        sayable_words = self._get_sayable_words(proc_sentence, self._word_dict)
        return self._get_sayable_sentence_arr(proc_sentence, sayable_words)

    def get_sentence_string(self, sentence: str) -> str:
        """Converts input sentence string to generalized output sentence string.
        Ex: "hello, there" -> "hello , there"

        Args:
            sentence (str): Sentence string.

        Returns:
            str: Generalized sentence string.
        """
        proc_sentence = self._process_sentence(sentence)
        sayable_words = self._get_sayable_words(
            proc_sentence, self._word_dict)
        sayable_sent_arr = self._get_sayable_sentence_arr(
            proc_sentence, sayable_words)

        return self._create_sentence_string(sayable_sent_arr)

    def get_generated_sentences(self) -> List[str]:
        """Gets the previously generated sentence strings

        Returns:
            List[str]: List of sentence strings generated previously
        """
        return [entry["sentence"] for entry in self._db.all()]

    def generate_audio(self, sentence: str, dry_run=False) -> Sentence:
        """Generates an audio file from the given sentence

        Args:
            sentence (string): Sentence string to be generated
            dry_run (bool, optional): Don't generate audio. Defaults to False.
        Returns:
            Sentence: Information about generated sentence.
        """

        log.info(f"Asked to generate {sentence}")
        proc_sentence = self._process_sentence(sentence)

        sayable_words = self._get_sayable_words(proc_sentence, self._word_dict)
        sayable_sent_arr = self._get_sayable_sentence_arr(
            proc_sentence, sayable_words)
        sayable_sent_str = self._create_sentence_string(sayable_sent_arr)
        unsayable_worlds = self._get_unsayable_words(
            proc_sentence, sayable_words)

        sentence_info = Sentence(
            sentence=sayable_sent_str,
            sayable=sayable_words,
            unsayable=unsayable_worlds,
            audio=None,
        )

        if dry_run:
            return sentence_info

        log.debug(f"Generating {sayable_sent_str}")

        words_audio = []

        # Only create sentence if there are words to put in it
        if len(sayable_words) == 0:
            log.warning(
                f"Can't say any words in {sentence}, not generating")
            return sentence_info

        # If sentence doesn't exist, create it (or overwrite if we are forcing remakes)
        sent_id = self._db.insert({"sentence": sayable_sent_str})

        for word in sayable_sent_arr:
            if word in self._punctuation_timing:
                words_audio.append(AudioSegment.silent(
                    self._punctuation_timing[word]))
            else:
                words_audio.append(AudioSegment.from_file(
                    self._word_dict[word], self._audio_format))

        sentence_audio = words_audio.pop(0)
        for word_audio in words_audio:
            word_audio = word_audio.set_frame_rate(sentence_audio.frame_rate)
            sentence_audio = sentence_audio + word_audio
        sentence_info.audio = sentence_audio

        return sentence_info

    def get_generated_sentences_dict(self) -> Dict[int, str]:
        """Gets the previously generated sentence strings
        along with their corresponding ID in the database

        Returns:
            Dict[int, str]: Dict of sentence and id pairs
        """
        entries = self._db.all()
        return {entries.index(entry): entry["sentence"] for entry in entries}


if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG)
    parser = argparse.ArgumentParser(
        description='Generate a sentence using a voice')
    parser.add_argument('-s', '--voice-dir', type=str, required=True, help='Path to folder with voice audio files')
    parser.add_argument('-d', '--db', type=str, required=True, help='Path to store database file')
    parser.add_argument('-f', '--format', type=str, required=False, default='wav', help='Audio format to export as')
    parser.add_argument('sentence', type=str)
    args = parser.parse_args()

    voice_dir = Path(args.voice_dir)
    if not voice_dir.is_dir():
        print(f'Voice dir at {voice_dir} does not exist!')
        sys.exit(1)

    db_path = Path(args.db)

    voice = Voice(path=voice_dir, db_path=db_path)
    sentence = voice.generate_audio(args.sentence)
    if sentence is None or sentence.audio is None:
        sys.exit(1)

    output_path = Path.cwd().joinpath(f"{'-'.join(sentence.sayable)}.{args.format}")

    sentence.audio.export(output_path, format=args.format)
