import json
import pathlib
from typing import Dict, List, Union

from utils.hparams import hparams

PAD_INDEX = 0


class PhonemeDictionary:
    def __init__(
            self,
            dictionaries: Dict[str, pathlib.Path],
            extra_phonemes: List[str] = None,
            merged_groups: List[List[str]] = None
    ):
        # Step 1: Collect all phonemes
        all_phonemes = {'AP', 'SP'}
        if extra_phonemes:
            for ph in extra_phonemes:
                if '/' in ph:
                    lang, name = ph.split('/', maxsplit=1)
                    if lang not in dictionaries:
                        raise ValueError(
                            f"Invalid phoneme tag '{ph}' in extra phonemes: "
                            f"unrecognized language name '{lang}'."
                        )
                    if name in all_phonemes:
                        raise ValueError(
                            f"Invalid phoneme tag '{ph}' in extra phonemes: "
                            f"short name conflicts with existing tag."
                        )
                all_phonemes.add(ph)
        self._multi_langs = len(dictionaries) > 1
        for lang, dict_path in dictionaries.items():
            with open(dict_path, 'r', encoding='utf8') as dict_file:
                for line in dict_file:
                    _, phonemes = line.strip().split('\t')
                    phonemes = phonemes.split()
                    for phoneme in phonemes:
                        if '/' in phoneme:
                            raise ValueError(
                                f"Invalid phoneme tag '{phoneme}' in dictionary '{dict_path}': "
                                f"should not contain the reserved character '/'."
                            )
                        if phoneme in all_phonemes:
                            continue
                        if self._multi_langs:
                            all_phonemes.add(f'{lang}/{phoneme}')
                        else:
                            all_phonemes.add(phoneme)
        # Step 2: Parse merged phoneme groups
        if merged_groups is None:
            merged_groups = []
        else:
            _merged_groups = []
            for group in merged_groups:
                _group = []
                for phoneme in group:
                    if '/' in phoneme:
                        lang, name = phoneme.split('/', maxsplit=1)
                        if lang not in dictionaries:
                            raise ValueError(
                                f"Invalid phoneme tag '{phoneme}' in merged group: "
                                f"unrecognized language name '{lang}'."
                            )
                        if self._multi_langs:
                            element = phoneme
                        else:
                            element = name
                    else:
                        element = phoneme
                    if element not in all_phonemes:
                        raise ValueError(
                            f"Invalid phoneme tag '{phoneme}' in merged group: "
                            f"not found in phoneme set."
                        )
                    _group.append(element)
                _merged_groups.append(_group)
            merged_groups = [set(phones) for phones in _merged_groups if len(phones) > 1]
        # Step 3: Build phoneme index
        merged_phonemes_inverted_index = {}
        for idx, group in enumerate(merged_groups):
            other_idx = None
            for phoneme in group:
                if phoneme in merged_phonemes_inverted_index:
                    other_idx = merged_phonemes_inverted_index[phoneme]
                    break
            target_idx = idx if other_idx is None else other_idx
            for phoneme in group:
                merged_phonemes_inverted_index[phoneme] = target_idx
            if other_idx is not None:
                merged_groups[other_idx] |= group
                group.clear()
        phone_to_id = {}
        id_to_phone = []
        cross_lingual_phonemes = set()
        idx = 1
        for phoneme in sorted(all_phonemes):
            if phoneme in merged_phonemes_inverted_index:
                has_assigned = True
                for alias in merged_groups[merged_phonemes_inverted_index[phoneme]]:
                    if alias not in phone_to_id:
                        has_assigned = False
                        phone_to_id[alias] = idx
                if not has_assigned:
                    merged_group = sorted(merged_groups[merged_phonemes_inverted_index[phoneme]])
                    merged_from_langs = {
                        (alias.split('/', maxsplit=1)[0] if '/' in alias else None)
                        for alias in merged_group
                    }
                    id_to_phone.append(tuple(merged_group))
                    idx += 1
                    if len(merged_from_langs) > 1:
                        cross_lingual_phonemes.update(ph for ph in merged_group if '/' in ph)
            else:
                phone_to_id[phoneme] = idx
                id_to_phone.append(phoneme)
                idx += 1
        self._phone_to_id: Dict[str, int] = phone_to_id
        self._id_to_phone: List[Union[str, tuple]] = id_to_phone
        self._cross_lingual_phonemes = frozenset(cross_lingual_phonemes)

    @property
    def vocab_size(self):
        return len(self._id_to_phone) + 1

    def __len__(self):
        return self.vocab_size

    @property
    def cross_lingual_phonemes(self):
        return self._cross_lingual_phonemes

    def is_cross_lingual(self, phone):
        return phone in self._cross_lingual_phonemes

    def encode_one(self, phone, lang=None):
        if '/' in phone:
            lang, phone = phone.split('/', maxsplit=1)
        if lang is None or not self._multi_langs or phone in self._phone_to_id:
            return self._phone_to_id[phone]
        if '/' not in phone:
            phone = f'{lang}/{phone}'
        return self._phone_to_id[phone]

    def encode(self, sentence, lang=None):
        phones = sentence.strip().split() if isinstance(sentence, str) else sentence
        return [self.encode_one(phone, lang=lang) for phone in phones]

    def decode_one(self, idx, lang=None, scalar=True):
        if idx <= 0:
            return None
        phone = self._id_to_phone[idx - 1]
        if not scalar or isinstance(phone, str):
            return phone
        if lang is None or not self._multi_langs:
            return phone[0]
        for alias in phone:
            if alias.startswith(f'{lang}/'):
                return alias
        return phone[0]

    def decode(self, ids, lang=None, scalar=True):
        ids = list(ids)
        return ' '.join([
            self.decode_one(i, lang=lang, scalar=scalar)
            for i in ids
            if i >= 1
        ])

    def dump(self, filename):
        with open(filename, 'w', encoding='utf8') as fp:
            json.dump(self._phone_to_id, fp, ensure_ascii=False, indent=2)


_dictionary = None


def load_phoneme_dictionary() -> PhonemeDictionary:
    if _dictionary is not None:
        return _dictionary
    config_dicts = hparams.get('dictionaries')
    if config_dicts is not None:
        dicts = {}
        for lang, config_dict_path in config_dicts.items():
            dict_path = pathlib.Path(hparams['work_dir']) / f'dictionary-{lang}.txt'
            if not dict_path.exists():
                dict_path = pathlib.Path(config_dict_path)
            if not dict_path.exists():
                raise FileNotFoundError(
                    f"Could not locate dictionary for language '{lang}'."
                )
            dicts[lang] = dict_path
    else:
        dict_path = pathlib.Path(hparams['work_dir']) / 'dictionary.txt'
        if not dict_path.exists():
            dict_path = pathlib.Path(hparams['dictionary'])
        if not dict_path.exists():
            raise FileNotFoundError(
                f"Could not locate dictionary file."
            )
        dicts = {
            'default': dict_path
        }
    return PhonemeDictionary(
        dictionaries=dicts,
        extra_phonemes=hparams.get('extra_phonemes'),
        merged_groups=hparams.get('merged_phoneme_groups')
    )
