#!/bin/bash

export DEBIAN_FRONTEND=noninteractive
export LOCK_FILE="/var/lib/dpkg/lock-frontend"
export LOCK_FILE2="/var/lib/apt/lists/lock"
export WAIT_INTERVAL=5
export MAX_WAIT=300
export ELAPSED=0

sudo systemctl stop apt-daily.timer
sudo systemctl stop apt-daily.service
sudo systemctl stop apt-daily-upgrade.timer
sudo systemctl stop apt-daily-upgrade.service
sudo systemctl stop unattended-upgrades.service
sudo systemctl stop packagekit

while fuser $LOCK_FILE $LOCK_FILE2 >/dev/null 2>&1; do
    sleep $WAIT_INTERVAL
    ELAPSED=$((ELAPSED + WAIT_INTERVAL))
    if [ $ELAPSED -ge $MAX_WAIT ]; then
        PIDS=$(fuser $LOCK_FILE $LOCK_FILE2 2>/dev/null)
        if [ -n "$PIDS" ]; then
            kill -9 $PIDS
        fi
        DEBIAN_FRONTEND=noninteractive dpkg --configure -a
        break
    fi
done

PIDS=$(pgrep -f "apt|dpkg|unattended|packagekit" | grep -v $$)
if [ -n "$PIDS" ]; then
    echo "Killing processes: $PIDS"
    kill -9 $PIDS 2>/dev/null
fi

sleep 3

rm -f /var/lib/apt/lists/lock
rm -f /var/cache/apt/archives/lock
rm -f /var/lib/dpkg/lock*
DEBIAN_FRONTEND=noninteractive dpkg --configure -a

DEBIAN_FRONTEND=noninteractive apt update -y


git clone https://huggingface.co/spaces/hexgrad/Kokoro-TTS /root/Kokoro-TTS

DEBIAN_FRONTEND=noninteractive apt install python3.10-venv espeak-ng supervisor ffmpeg net-tools -y

python3 -m venv /root/Kokoro-TTS/venv

/root/Kokoro-TTS/venv/bin/pip install --upgrade pip setuptools wheel

/root/Kokoro-TTS/venv/bin/pip install soundfile
/root/Kokoro-TTS/venv/bin/pip install spaces

/root/Kokoro-TTS/venv/bin/pip install git+https://github.com/hexgrad/kokoro.git

cat << 'EOF' > /root/Kokoro-TTS/app.py
import spaces
from kokoro import KModel, KPipeline
import gradio as gr
import os
import random
import torch

IS_DUPLICATE = not os.getenv('SPACE_ID', '').startswith('hexgrad/')
CUDA_AVAILABLE = torch.cuda.is_available()
if not IS_DUPLICATE:
    import kokoro
    import misaki
    print('DEBUG', kokoro.__version__, CUDA_AVAILABLE, misaki.__version__)

CHAR_LIMIT = None

models = {gpu: KModel().to('cuda' if gpu else 'cpu').eval() for gpu in [False] + ([True] if CUDA_AVAILABLE else [])}
pipelines = {lang_code: KPipeline(lang_code=lang_code, model=False) for lang_code in 'ab'}
pipelines['a'].g2p.lexicon.golds['kokoro'] = 'kˈOkəɹO'
pipelines['b'].g2p.lexicon.golds['kokoro'] = 'kˈQkəɹQ'

@spaces.GPU(duration=3600)
def forward_gpu(ps, ref_s, speed):
    return models[True](ps, ref_s, speed)

def generate_all(text, voice='af_heart', speed=1, use_gpu=CUDA_AVAILABLE):
    text = text if CHAR_LIMIT is None else text.strip()[:CHAR_LIMIT]
    pipeline = pipelines[voice[0]]
    pack = pipeline.load_voice(voice)
    use_gpu = use_gpu and CUDA_AVAILABLE

    for _, ps, _ in pipeline(text, voice, speed):
        ref_s = pack[len(ps)-1]
        try:
            if use_gpu:
                audio = forward_gpu(ps, ref_s, speed)
            else:
                audio = models[False](ps, ref_s, speed)
        except gr.exceptions.Error as e:
            if use_gpu:
                gr.Warning(str(e))
                gr.Info('Switching to CPU')
                audio = models[False](ps, ref_s, speed)
            else:
                raise gr.Error(e)
        yield 24000, audio.numpy()

def get_random_quote():
    with open('en.txt', 'r') as r:
        random_quotes = [line.strip() for line in r]
    return random.choice(random_quotes)

def get_gatsby():
    with open('gatsby5k.md', 'r') as r:
        return r.read().strip()

def get_frankenstein():
    with open('frankenstein5k.md', 'r') as r:
        return r.read().strip()

CHOICES = {
    '🇺🇸 🚺 Heart ❤️': 'af_heart',
    '🇺🇸 🚺 Bella 🔥': 'af_bella',
    '🇺🇸 🚺 Nicole 🎧': 'af_nicole',
    '🇺🇸 🚺 Aoede': 'af_aoede',
    '🇺🇸 🚺 Kore': 'af_kore',
    '🇺🇸 🚺 Sarah': 'af_sarah',
    '🇺🇸 🚺 Nova': 'af_nova',
    '🇺🇸 🚺 Sky': 'af_sky',
    '🇺🇸 🚺 Alloy': 'af_alloy',
    '🇺🇸 🚺 Jessica': 'af_jessica',
    '🇺🇸 🚺 River': 'af_river',
    '🇺🇸 🚹 Michael': 'am_michael',
    '🇺🇸 🚹 Fenrir': 'am_fenrir',
    '🇺🇸 🚹 Puck': 'am_puck',
    '🇺🇸 🚹 Echo': 'am_echo',
    '🇺🇸 🚹 Eric': 'am_eric',
    '🇺🇸 🚹 Liam': 'am_liam',
    '🇺🇸 🚹 Onyx': 'am_onyx',
    '🇺🇸 🚹 Santa': 'am_santa',
    '🇺🇸 🚹 Adam': 'am_adam',
    '🇬🇧 🚺 Emma': 'bf_emma',
    '🇬🇧 🚺 Isabella': 'bf_isabella',
    '🇬🇧 🚺 Alice': 'bf_alice',
    '🇬🇧 🚺 Lily': 'bf_lily',
    '🇬🇧 🚹 George': 'bm_george',
    '🇬🇧 🚹 Fable': 'bm_fable',
    '🇬🇧 🚹 Lewis': 'bm_lewis',
    '🇬🇧 🚹 Daniel': 'bm_daniel',
}
for v in CHOICES.values():
    pipelines[v[0]].load_voice(v)

with gr.Blocks() as app:
    text = gr.Textbox(
        label='Input Text',
        info=f"Up to ~500 characters per Generate, or {'∞' if CHAR_LIMIT is None else CHAR_LIMIT} characters per Stream"
    )
    voice = gr.Dropdown(list(CHOICES.items()), value='af_heart', label='Voice')
    use_gpu = gr.Dropdown([('GPU 🚀', True), ('CPU 🐌', False)], value=CUDA_AVAILABLE, label='Hardware')
    speed = gr.Slider(minimum=0.5, maximum=2, value=1, step=0.1, label='Speed')
    out_audio = gr.Audio(label='Output Audio Stream', interactive=False, streaming=True, autoplay=True)
    generate_btn = gr.Button('Generate', variant='primary')

    generate_btn.click(fn=generate_all, inputs=[text, voice, speed, use_gpu], outputs=[out_audio])

if __name__ == '__main__':
    app.queue(api_open=True).launch(share=True, server_name="0.0.0.0", server_port=7860)
EOF

echo "[program:kokoro-tts]
command=/root/Kokoro-TTS/venv/bin/python /root/Kokoro-TTS/app.py
directory=/root/Kokoro-TTS
autostart=true
autorestart=true
startretries=3
user=root
stdout_logfile=/var/log/kokoro-tts.log
stderr_logfile=/var/log/kokoro-tts.err
environment=PYTHONUNBUFFERED=\"1\"
" > /etc/supervisor/conf.d/kokoro-tts.conf

supervisorctl reread
supervisorctl update
supervisorctl start kokoro-tts


START_TIME=$(date +%s)
TIMEOUT=$((10 * 60))  # 10 minutes in seconds

while true; do
    NOW=$(date +%s)
    ELAPSED=$((NOW - START_TIME))

    if (( ELAPSED >= TIMEOUT )); then
        echo "$(date '+%F %T') Timeout reached after 10 minutes. Exiting."
        break
    fi

    if curl -s --max-time 5 http://127.0.0.1:7860 >/dev/null; then
        echo "$(date '+%F %T') Kokoro TTS service is working. Exiting."
        break
    else
        echo "$(date '+%F %T') Kokoro TTS service is not reachable. Checking again in 10 seconds."
        sleep 10
    fi
done

