# Copyright 2024 Jordi Mas i Hernàndez <jmas@softcatala.org>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging

from typing import List

import numpy as np
import scipy.io.wavfile
import torch

from transformers import AutoTokenizer, VitsModel

from open_dubbing.text_to_speech import TextToSpeech, Voice


class TextToSpeechMMS(TextToSpeech):

    def __init__(self, device="cpu"):
        super().__init__()
        self.device = device
        logging.getLogger("transformers").setLevel(logging.ERROR)

    def get_available_voices(self, language_code: str) -> List[Voice]:
        return []

    def _convert_text_to_speech(
        self,
        *,
        assigned_voice: str,
        target_language: str,
        output_filename: str,
        text: str,
        speed: float,
    ) -> str:

        logging.debug(f"TextToSpeechMMS._convert_text_to_speech: {text}")
        local_files_only = False

        # Load pre-trained model and tokenizer
        model = VitsModel.from_pretrained(
            f"facebook/mms-tts-{target_language}", local_files_only=local_files_only
        ).to(self.device)
        tokenizer = AutoTokenizer.from_pretrained(
            f"facebook/mms-tts-{target_language}", local_files_only=local_files_only
        )
        inputs = tokenizer(text, return_tensors="pt").to(self.device)

        # Generate waveform
        with torch.no_grad():
            output = model(**inputs).waveform

        # Convert waveform to NumPy array and scale to 16-bit PCM
        # Assuming `output` is a 2D tensor with shape (batch_size, samples)
        output_np = output.squeeze().cpu().numpy()  # Remove batch dimension if present
        output_np = np.clip(output_np, -1, 1)  # Clip values to be between -1 and 1
        output_np = (output_np * 32767).astype(np.int16)  # Scale to 16-bit PCM

        # Get the sampling rate
        sampling_rate = model.config.sampling_rate

        # Write to WAV file
        wav_file = output_filename.replace(".mp3", ".wav")
        scipy.io.wavfile.write(wav_file, rate=sampling_rate, data=output_np)

        self._convert_to_mp3(wav_file, output_filename)
        logging.debug(
            f"text_to_speech.client.synthesize_speech: output_filename: '{output_filename}'"
        )
        return output_filename

    # Reference: https://dl.fbaipublicfiles.com/mms/tts/all-tts-languages.html
    def get_languages(self):
        return [
            "abi",
            "abp",
            "aca",
            "acd",
            "ace",
            "acf",
            "ach",
            "acn",
            "acr",
            "acu",
            "ade",
            "adh",
            "adj",
            "adx",
            "aeu",
            "agd",
            "agg",
            "agn",
            "agr",
            "agu",
            "agx",
            "aha",
            "ahk",
            "aia",
            "aka",
            "akb",
            "ake",
            "akp",
            "alj",
            "alp",
            "alt",
            "alz",
            "ame",
            "amf",
            "amh",
            "ami",
            "amk",
            "ann",
            "any",
            "aoz",
            "apb",
            "apr",
            "ara",
            "arl",
            "asa",
            "asg",
            "asm",
            "ata",
            "atb",
            "atg",
            "ati",
            "atq",
            "ava",
            "avn",
            "avu",
            "awa",
            "awb",
            "ayo",
            "ayr",
            "ayz",
            "azb",
            "azg",
            "azz",
            "bak",
            "bam",
            "ban",
            "bao",
            "bav",
            "bba",
            "bbb",
            "bbc",
            "bbo",
            "bcl",
            "bcw",
            "bdg",
            "bdh",
            "bdq",
            "bdu",
            "bdv",
            "beh",
            "bem",
            "ben",
            "bep",
            "bex",
            "bfa",
            "bfo",
            "bfy",
            "bfz",
            "bgc",
            "bgq",
            "bgr",
            "bgt",
            "bgw",
            "bha",
            "bht",
            "bhz",
            "bib",
            "bim",
            "bis",
            "biv",
            "bjr",
            "bjv",
            "bjw",
            "bjz",
            "bkd",
            "bkv",
            "blh",
            "blt",
            "blx",
            "blz",
            "bmq",
            "bmr",
            "bmu",
            "bmv",
            "bng",
            "bno",
            "bnp",
            "boa",
            "bod",
            "boj",
            "bom",
            "bor",
            "bov",
            "box",
            "bpr",
            "bps",
            "bqc",
            "bqi",
            "bqj",
            "bqp",
            "bru",
            "bsc",
            "bsq",
            "bss",
            "btd",
            "bts",
            "btt",
            "btx",
            "bud",
            "bul",
            "bus",
            "bvc",
            "bvz",
            "bwq",
            "bwu",
            "byr",
            "bzh",
            "bzi",
            "bzj",
            "caa",
            "cab",
            "cap",
            "car",
            "cas",
            "cat",
            "cax",
            "cbc",
            "cbi",
            "cbr",
            "cbs",
            "cbt",
            "cbu",
            "cbv",
            "cce",
            "cco",
            "cdj",
            "ceb",
            "ceg",
            "cek",
            "cfm",
            "cgc",
            "che",
            "chf",
            "chv",
            "chz",
            "cjo",
            "cjp",
            "cjs",
            "cko",
            "ckt",
            "cla",
            "cle",
            "cly",
            "cme",
            "cmr",
            "cnh",
            "cni",
            "cnl",
            "cnt",
            "coe",
            "cof",
            "cok",
            "con",
            "cot",
            "cou",
            "cpa",
            "cpb",
            "cpu",
            "crh",
            "crn",
            "crq",
            "crs",
            "crt",
            "csk",
            "cso",
            "ctd",
            "ctg",
            "cto",
            "ctu",
            "cuc",
            "cui",
            "cuk",
            "cul",
            "cwa",
            "cwe",
            "cwt",
            "cya",
            "cym",
            "daa",
            "dah",
            "dar",
            "dbj",
            "dbq",
            "ddn",
            "ded",
            "des",
            "deu",
            "dga",
            "dgi",
            "dgk",
            "dgo",
            "dgr",
            "dhi",
            "did",
            "dig",
            "dik",
            "dip",
            "div",
            "djk",
            "dnt",
            "dnw",
            "dop",
            "dos",
            "dsh",
            "dso",
            "dtp",
            "dts",
            "dug",
            "dwr",
            "dyi",
            "dyo",
            "dyu",
            "dzo",
            "eip",
            "eka",
            "ell",
            "emp",
            "enb",
            "eng",
            "enx",
            "ese",
            "ess",
            "eus",
            "evn",
            "ewe",
            "eza",
            "fal",
            "fao",
            "far",
            "fas",
            "fij",
            "fin",
            "flr",
            "fmu",
            "fon",
            "fra",
            "frd",
            "ful",
            "gai",
            "gam",
            "gau",
            "gbi",
            "gbk",
            "gbm",
            "gbo",
            "gde",
            "geb",
            "gej",
            "gil",
            "gjn",
            "gkn",
            "gld",
            "glk",
            "gmv",
            "gna",
            "gnd",
            "gng",
            "gog",
            "gor",
            "gqr",
            "grc",
            "gri",
            "grn",
            "grt",
            "gso",
            "gub",
            "guc",
            "gud",
            "guh",
            "guj",
            "guk",
            "gum",
            "guo",
            "guq",
            "guu",
            "gux",
            "gvc",
            "gvl",
            "gwi",
            "gwr",
            "gym",
            "gyr",
            "had",
            "hag",
            "hak",
            "hap",
            "hat",
            "hau",
            "hay",
            "heb",
            "heh",
            "hif",
            "hig",
            "hil",
            "hin",
            "hlb",
            "hlt",
            "hne",
            "hnn",
            "hns",
            "hoc",
            "hoy",
            "hto",
            "hub",
            "hui",
            "hun",
            "huu",
            "huv",
            "hvn",
            "hwc",
            "hyw",
            "iba",
            "icr",
            "idd",
            "ifa",
            "ifb",
            "ife",
            "ifk",
            "ifu",
            "ify",
            "ign",
            "ikk",
            "ilb",
            "ilo",
            "imo",
            "inb",
            "ind",
            "iou",
            "ipi",
            "iqw",
            "iri",
            "irk",
            "isl",
            "itl",
            "itv",
            "izr",
            "izz",
            "jac",
            "jam",
            "jav",
            "jbu",
            "jen",
            "jic",
            "jiv",
            "jmc",
            "jmd",
            "jun",
            "juy",
            "jvn",
            "kaa",
            "kab",
            "kac",
            "kak",
            "kan",
            "kao",
            "kaq",
            "kay",
            "kaz",
            "kbo",
            "kbp",
            "kbq",
            "kbr",
            "kby",
            "kca",
            "kcg",
            "kdc",
            "kde",
            "kdh",
            "kdi",
            "kdj",
            "kdl",
            "kdn",
            "kdt",
            "kek",
            "ken",
            "keo",
            "ker",
            "key",
            "kez",
            "kfb",
            "kfw",
            "kfx",
            "khg",
            "khm",
            "khq",
            "kia",
            "kij",
            "kik",
            "kin",
            "kir",
            "kjb",
            "kje",
            "kjg",
            "kjh",
            "kki",
            "kkj",
            "kle",
            "klu",
            "klv",
            "klw",
            "kma",
            "kmd",
            "kml",
            "kmu",
            "knb",
            "kne",
            "knf",
            "knj",
            "knk",
            "kno",
            "kog",
            "kor",
            "kpq",
            "kps",
            "kpv",
            "kpy",
            "kpz",
            "kqe",
            "kqp",
            "kqr",
            "kqy",
            "krc",
            "kri",
            "krj",
            "krl",
            "krr",
            "kru",
            "ksb",
            "ksr",
            "kss",
            "ktb",
            "ktj",
            "kub",
            "kue",
            "kum",
            "kus",
            "kvn",
            "kvw",
            "kwd",
            "kwf",
            "kwi",
            "kxc",
            "kxf",
            "kxm",
            "kxv",
            "kyb",
            "kyc",
            "kyf",
            "kyg",
            "kyo",
            "kyq",
            "kyu",
            "kyz",
            "kzf",
            "lac",
            "laj",
            "lam",
            "lao",
            "las",
            "lat",
            "lav",
            "law",
            "lbj",
            "lbw",
            "lcp",
            "lee",
            "lef",
            "lem",
            "lew",
            "lex",
            "lgg",
            "lgl",
            "lhu",
            "lia",
            "lid",
            "lif",
            "lip",
            "lis",
            "lje",
            "ljp",
            "llg",
            "lln",
            "lme",
            "lnd",
            "lns",
            "lob",
            "lok",
            "lom",
            "lon",
            "loq",
            "lsi",
            "lsm",
            "luc",
            "lug",
            "lwo",
            "lww",
            "lzz",
            "mad",
            "mag",
            "mah",
            "mai",
            "maj",
            "mak",
            "mal",
            "maq",
            "mar",
            "maw",
            "maz",
            "mbb",
            "mbc",
            "mbh",
            "mbj",
            "mbt",
            "mbu",
            "mbz",
            "mca",
            "mcb",
            "mcd",
            "mco",
            "mcp",
            "mcq",
            "mcu",
            "mda",
            "mdv",
            "mdy",
            "med",
            "mee",
            "mej",
            "men",
            "meq",
            "met",
            "mev",
            "mfe",
            "mfh",
            "mfi",
            "mfk",
            "mfq",
            "mfy",
            "mfz",
            "mgd",
            "mge",
            "mgh",
            "mgo",
            "mhi",
            "mhr",
            "mhu",
            "mhx",
            "mhy",
            "mib",
            "mie",
            "mif",
            "mih",
            "mil",
            "mim",
            "min",
            "mio",
            "mip",
            "miq",
            "mit",
            "miy",
            "miz",
            "mjl",
            "mjv",
            "mkl",
            "mkn",
            "mlg",
            "mmg",
            "mnb",
            "mnf",
            "mnk",
            "mnw",
            "mnx",
            "moa",
            "mog",
            "mon",
            "mop",
            "mor",
            "mos",
            "mox",
            "moz",
            "mpg",
            "mpm",
            "mpp",
            "mpx",
            "mqb",
            "mqf",
            "mqj",
            "mqn",
            "mrw",
            "msy",
            "mtd",
            "mtj",
            "mto",
            "muh",
            "mup",
            "mur",
            "muv",
            "muy",
            "mvp",
            "mwq",
            "mwv",
            "mxb",
            "mxq",
            "mxt",
            "mxv",
            "mya",
            "myb",
            "myk",
            "myl",
            "myv",
            "myx",
            "myy",
            "mza",
            "mzi",
            "mzj",
            "mzk",
            "mzm",
            "mzw",
            "nab",
            "nag",
            "nan",
            "nas",
            "naw",
            "nca",
            "nch",
            "ncj",
            "ncl",
            "ncu",
            "ndj",
            "ndp",
            "ndv",
            "ndy",
            "ndz",
            "neb",
            "new",
            "nfa",
            "nfr",
            "nga",
            "ngl",
            "ngp",
            "ngu",
            "nhe",
            "nhi",
            "nhu",
            "nhw",
            "nhx",
            "nhy",
            "nia",
            "nij",
            "nim",
            "nin",
            "nko",
            "nlc",
            "nld",
            "nlg",
            "nlk",
            "nmz",
            "nnb",
            "nnq",
            "nnw",
            "noa",
            "nod",
            "nog",
            "not",
            "npl",
            "npy",
            "nst",
            "nsu",
            "ntm",
            "ntr",
            "nuj",
            "nus",
            "nuz",
            "nwb",
            "nxq",
            "nya",
            "nyf",
            "nyn",
            "nyo",
            "nyy",
            "nzi",
            "obo",
            "oku",
            "old",
            "omw",
            "onb",
            "ood",
            "orm",
            "ory",
            "oss",
            "ote",
            "otq",
            "ozm",
            "pab",
            "pad",
            "pag",
            "pam",
            "pan",
            "pao",
            "pap",
            "pau",
            "pbb",
            "pbc",
            "pbi",
            "pce",
            "pcm",
            "peg",
            "pez",
            "pib",
            "pil",
            "pir",
            "pis",
            "pjt",
            "pkb",
            "pls",
            "plw",
            "pmf",
            "pny",
            "poi",
            "pol",
            "por",
            "poy",
            "ppk",
            "pps",
            "prf",
            "prk",
            "prt",
            "pse",
            "pss",
            "ptu",
            "pui",
            "pwg",
            "pww",
            "pxm",
            "qub",
            "quf",
            "quh",
            "qul",
            "quw",
            "quy",
            "quz",
            "qvc",
            "qve",
            "qvh",
            "qvm",
            "qvn",
            "qvo",
            "qvs",
            "qvw",
            "qvz",
            "qwh",
            "qxh",
            "qxl",
            "qxn",
            "qxo",
            "qxr",
            "rah",
            "rai",
            "rap",
            "rav",
            "raw",
            "rej",
            "rel",
            "rgu",
            "rhg",
            "ril",
            "rim",
            "rjs",
            "rkt",
            "rmo",
            "rng",
            "rnl",
            "rol",
            "ron",
            "rop",
            "rro",
            "rub",
            "ruf",
            "rug",
            "run",
            "rus",
            "sab",
            "sag",
            "sah",
            "saj",
            "saq",
            "sas",
            "sba",
            "sbd",
            "sbl",
            "sbp",
            "sch",
            "sck",
            "sda",
            "sea",
            "seh",
            "ses",
            "sey",
            "sgb",
            "sgj",
            "sgw",
            "shi",
            "shk",
            "shn",
            "sho",
            "shp",
            "sid",
            "sig",
            "sil",
            "sja",
            "sjm",
            "sld",
            "slu",
            "sml",
            "smo",
            "sna",
            "sne",
            "snn",
            "snp",
            "snw",
            "som",
            "soy",
            "spa",
            "spp",
            "spy",
            "sqi",
            "sri",
            "srm",
            "srn",
            "srx",
            "stn",
            "stp",
            "suc",
            "suk",
            "sun",
            "sur",
            "sus",
            "suv",
            "suz",
            "swe",
            "swh",
            "sxb",
            "sxn",
            "sya",
            "syl",
            "sza",
            "tac",
            "taj",
            "tam",
            "tao",
            "tap",
            "taq",
            "tat",
            "tav",
            "tbc",
            "tbg",
            "tbk",
            "tbl",
            "tby",
            "tbz",
            "tca",
            "tcc",
            "tcs",
            "tcz",
            "tdj",
            "ted",
            "tee",
            "tel",
            "tem",
            "teo",
            "ter",
            "tes",
            "tew",
            "tex",
            "tfr",
            "tgj",
            "tgk",
            "tgl",
            "tgo",
            "tgp",
            "tha",
            "thk",
            "thl",
            "tih",
            "tik",
            "tir",
            "tkr",
            "tlb",
            "tlj",
            "tly",
            "tmc",
            "tmf",
            "tna",
            "tng",
            "tnk",
            "tnn",
            "tnp",
            "tnr",
            "tnt",
            "tob",
            "toc",
            "toh",
            "tom",
            "tos",
            "tpi",
            "tpm",
            "tpp",
            "tpt",
            "trc",
            "tri",
            "trn",
            "trs",
            "tso",
            "tsz",
            "ttc",
            "tte",
            "tue",
            "tuf",
            "tuo",
            "tur",
            "tvw",
            "twb",
            "twe",
            "twu",
            "txa",
            "txq",
            "txu",
            "tye",
            "udm",
            "udu",
            "ukr",
            "unr",
            "upv",
            "ura",
            "urb",
            "urk",
            "urt",
            "ury",
            "usp",
            "vag",
            "vid",
            "vie",
            "vif",
            "vmw",
            "vmy",
            "vun",
            "vut",
            "wap",
            "war",
            "waw",
            "way",
            "wba",
            "wlo",
            "wlx",
            "wmw",
            "wob",
            "wsg",
            "wwa",
            "xal",
            "xdy",
            "xed",
            "xer",
            "xmm",
            "xnj",
            "xnr",
            "xog",
            "xon",
            "xrb",
            "xsb",
            "xsm",
            "xsr",
            "xsu",
            "xta",
            "xtd",
            "xte",
            "xtm",
            "xtn",
            "xua",
            "xuo",
            "yaa",
            "yad",
            "yal",
            "yam",
            "yao",
            "yas",
            "yat",
            "yaz",
            "yba",
            "ybb",
            "ycl",
            "ycn",
            "yea",
            "yka",
            "yli",
            "yor",
            "yre",
            "yua",
            "yuz",
            "yva",
            "zaa",
            "zab",
            "zac",
            "zad",
            "zae",
            "zai",
            "zam",
            "zao",
            "zaq",
            "zar",
            "zas",
            "zav",
            "zaw",
            "zca",
            "zga",
            "zim",
            "ziw",
            "zlm",
            "zmz",
            "zne",
            "zos",
            "zpc",
            "zpg",
            "zpi",
            "zpl",
            "zpm",
            "zpo",
            "zpt",
            "zpu",
            "zpz",
            "ztq",
            "zty",
            "zyb",
            "zyp",
            "zza",
        ]
