import gradio as gr
import torch
import numpy as np
import scipy.io.wavfile
from transformers import VitsModel, AutoTokenizer
import re

# Load model and tokenizer
model = VitsModel.from_pretrained("Somali-tts/somali_tts_model")
tokenizer = AutoTokenizer.from_pretrained("saleolow/somali-mms-tts")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()

# Numbers in Somali
number_words = {
    0: "eber", 1: "koow", 2: "labo", 3: "seddex", 4: "afar", 5: "shan",
    6: "lix", 7: "todobo", 8: "sideed", 9: "sagaal", 10: "toban",
    11: "toban iyo koow", 12: "toban iyo labo", 13: "toban iyo seddex",
    14: "toban iyo afar", 15: "toban iyo shan", 16: "toban iyo lix",
    17: "toban iyo todobo", 18: "toban iyo sideed", 19: "toban iyo sagaal",
    20: "labaatan", 30: "sodon", 40: "afartan", 50: "konton",
    60: "lixdan", 70: "todobaatan", 80: "sideetan", 90: "sagaashan",
    100: "boqol", 1000: "kun"
}

# Shortcut dictionary
shortcut_map = {
    "asc": "asalaamu caleykum",
    "wcs": "wacaleykum salaam",
    "fcn": "fiican",
    "xld": "xaaladda ka waran",
    "kwrn": "kawaran",
    "scw": "salalaahu caleyhi wa salam",
    "alx": "alxamdu lilaahi",
    "m.a": "maasha allah",
    "sthy": "side tahey",
    "sxp": "saaxiib"
}

# Countries dictionary (English to Somali)
country_map = {
    "somalia": "Soomaaliya",
    "ethiopia": "Itoobiya",
    "kenya": "Kenya",
    "djibouti": "Jabuuti",
    "sudan": "Suudaan",
    "Yeman": "yemaan",
    "uganda": "Ugaandha",
    "tanzania": "Tansaaniya",
    "egypt": "Masar",
    "libya": "Liibiya",
    "algeria": "Aljeeriya",
    "morocco": "Morooko",
    "tunisia": "Tuniisiya",
    "eritrea": "Eriteriya",
    "malawi": "Malaawi",
    "English": "ingiriis",
    "Spain": "isbeen",
    "Brazil": "baraasiil",
    "niger": "Niyjer",
    "Italy": "itaaliya",
    "united states": "Maraykanka",
    "china": "Shiinaha",
    "india": "Hindiya",
    "russia": "Ruushka",
    "Saudi Arabia": "Sucuudi Carabiya",
    "germany": "Jarmalka",
    "france": "Faransiiska",
    "japan": "Jabaan",
    "canada": "Kanada",
    "australia": "Australia"
}

def number_to_words(number):
    number = int(number)
    if number < 20:
        return number_words[number]
    elif number < 100:
        tens, unit = divmod(number, 10)
        return number_words[tens * 10] + (" iyo " + number_words[unit] if unit else "")
    elif number < 1000:
        hundreds, remainder = divmod(number, 100)
        part = (number_words[hundreds] + " boqol") if hundreds > 1 else "boqol"
        if remainder:
            part += " iyo " + number_to_words(remainder)
        return part
    elif number < 1000000:
        thousands, remainder = divmod(number, 1000)
        words = []
        if thousands == 1:
            words.append("kun")
        else:
            words.append(number_to_words(thousands) + " kun")
        if remainder >= 100:
            hundreds, rem2 = divmod(remainder, 100)
            if hundreds:
                boqol_text = (number_words[hundreds] + " boqol") if hundreds > 1 else "boqol"
                words.append(boqol_text)
            if rem2:
                words.append("iyo " + number_to_words(rem2))
        elif remainder:
            words.append("iyo " + number_to_words(remainder))
        return " ".join(words)
    elif number < 1000000000:
        millions, remainder = divmod(number, 1000000)
        words = []
        if millions == 1:
            words.append("milyan")
        else:
            words.append(number_to_words(millions) + " milyan")
        if remainder:
            words.append(number_to_words(remainder))
        return " ".join(words)
    else:
        return str(number)

def normalize_text(text):
    # ➤ Beddel ereyga zamzam dhammaan noocyadiisa (bilow, dhex, dhammaad)
    text = re.sub(r'(?i)(?<!\w)zamzam(?!\w)', 'samsam', text)

    # ➤ Bedel shortcuts - eray kasta oo qoraalka ku jira beddel
    def replace_shortcuts(match):
        word = match.group(0).lower()
        return shortcut_map.get(word, word)
    pattern = re.compile(r'\b(' + '|'.join(re.escape(k) for k in shortcut_map.keys()) + r')\b', re.IGNORECASE)
    text = pattern.sub(replace_shortcuts, text)

    # ➤ Bedel magacyada waddamada
    def replace_countries(match):
        word = match.group(0).lower()
        return country_map.get(word, word)
    country_pattern = re.compile(r'\b(' + '|'.join(re.escape(k) for k in country_map.keys()) + r')\b', re.IGNORECASE)
    text = country_pattern.sub(replace_countries, text)

    # ➤ Ka saar tirooyin leh koma iyo tobanle
    text = re.sub(r'(\d{1,3})(,\d{3})+', lambda m: m.group(0).replace(",", ""), text)
    text = re.sub(r'\.\d+', '', text)

    # ➤ Beddel tirooyin (1 → koow, iwm.)
    def replace_num(match):
        return number_to_words(match.group())
    text = re.sub(r'\d+', replace_num, text)

    # ➤ Beddel calaamadaha khaaska ah
    symbol_map = {
        '$': 'doolar',
        '=': 'egwal',
        '+': 'balaas',
        '#': 'haash'
    }
    for sym, word in symbol_map.items():
        text = text.replace(sym, ' ' + word + ' ')

    # ➤ Beddel xarfo khaas ah
    text = text.replace("KH", "qa").replace("Z", "S")
    text = text.replace("SH", "SHa'a").replace("DH", "Dha'a")

    # ➤ Haddii jumladda ku dhammaato zamzam ama samsam, ku dar " m" (silent m)
    if re.search(r'(?i)(zamzam|samsam)[\s\.\,\!\?]*$', text.strip()):
        text += " m"

    return text

def tts(text):
    paragraphs = text.strip().split("\n")
    audio_list = []
    max_chars = 500  # Qiyaasta ugu badan 2 daqiiqo
    warn_msg = ""

    for i, para in enumerate(paragraphs):
        para = para.strip()
        if not para:
            continue

        if len(para) > max_chars:
            warn_msg += f"❗ Qaybta {i+1} aad ayaa ka badan 2 daqiiqo. Waan kala jaray.\n"
            sub_parts = [para[j:j+max_chars] for j in range(0, len(para), max_chars)]
        else:
            sub_parts = [para]

        for part in sub_parts:
            norm_para = normalize_text(part)
            inputs = tokenizer(norm_para, return_tensors="pt").to(device)
            with torch.no_grad():
                waveform = model(**inputs).waveform.squeeze().cpu().numpy()

            pause = np.zeros(int(model.config.sampling_rate * 0.8))  # 0.8s pause
            audio_list.append(np.concatenate((waveform, pause)))

    final_audio = np.concatenate(audio_list)
    filename = "output.wav"
    scipy.io.wavfile.write(filename, rate=model.config.sampling_rate, data=(final_audio * 32767).astype(np.int16))

    if warn_msg:
        print(warn_msg)
    return filename

# Gradio interface
gr.Interface(
    fn=tts,
    inputs=gr.Textbox(label="Geli qoraal Soomaali ah", lines=10, placeholder="Ku qor 1 ama in ka badan paragraph..."),
    outputs=gr.Audio(label="Codka TTS"),
    title="Somali TTS",
    description="Ku qor qoraal Soomaaliyeed si aad u maqasho cod dabiici ah. Qoraalka ha ka badnaan 2 daqiiqo per jumlad."
).launch()