#!/bin/bash
set -e

MODELS_DIR="${MODELS_DIR:-/app/models}"
MODEL_BASE_URL="${MODEL_BASE_URL:-}"

# Model directories
PARAKEET_DIR="parakeet-tdt-0.6b-v3"
EOU_DIR="realtime_eou_120m-v1-onnx"
DIARIZATION_DIR="diarization"

download_from_url() {
    local name=$1
    local url=$2
    local dest="$MODELS_DIR/$name"

    if [ -d "$dest" ] && [ "$(ls -A $dest 2>/dev/null)" ]; then
        echo "Model $name already exists, skipping..."
        return 0
    fi

    echo "Downloading $name from URL..."
    mkdir -p "$dest"
    curl -L "$url" | tar -xz -C "$dest" --strip-components=1
    echo "Downloaded $name successfully"
}

download_from_hf() {
    local dest=$1
    local repo=$2

    if [ -d "$dest" ] && [ "$(ls -A $dest 2>/dev/null)" ]; then
        echo "Model $dest already exists, skipping..."
        return 0
    fi

    echo "Downloading from Hugging Face ($repo)..."
    mkdir -p "$dest"
    hf download "$repo" --local-dir "$dest"
    echo "Downloaded to $dest successfully"
}

download_from_hf_subdir() {
    local dest=$1
    local repo=$2
    local subdir=$3

    if [ -d "$dest" ] && [ "$(ls -A $dest 2>/dev/null)" ]; then
        echo "Model $dest already exists, skipping..."
        return 0
    fi

    echo "Downloading $subdir from Hugging Face ($repo)..."
    local tmpdir=$(mktemp -d)
    hf download "$repo" --include "$subdir/*" --local-dir "$tmpdir"

    # Move subdirectory contents to destination
    mkdir -p "$dest"
    mv "$tmpdir/$subdir"/* "$dest"/
    rm -rf "$tmpdir"
    echo "Downloaded to $dest successfully"
}

# Check if models exist
# TDT: encoder-model.onnx, encoder-model.onnx.data, decoder_joint-model.onnx, vocab.txt
# EOU: encoder.onnx, decoder_joint.onnx, tokenizer.json
# Diarization: diar_streaming_sortformer_4spk-v2.1.onnx
check_models_exist() {
    # TDT model files
    [ -f "$MODELS_DIR/$PARAKEET_DIR/encoder-model.onnx" ] && \
    [ -f "$MODELS_DIR/$PARAKEET_DIR/decoder_joint-model.onnx" ] && \
    [ -f "$MODELS_DIR/$PARAKEET_DIR/vocab.txt" ] && \
    # EOU model files
    [ -f "$MODELS_DIR/$EOU_DIR/encoder.onnx" ] && \
    [ -f "$MODELS_DIR/$EOU_DIR/decoder_joint.onnx" ] && \
    [ -f "$MODELS_DIR/$EOU_DIR/tokenizer.json" ] && \
    # Diarization model
    [ -f "$MODELS_DIR/$DIARIZATION_DIR/diar_streaming_sortformer_4spk-v2.1.onnx" ]
}

if check_models_exist; then
    echo "All models present"
else
    mkdir -p "$MODELS_DIR"

    if [ -n "$MODEL_BASE_URL" ]; then
        echo "Downloading models from custom URL..."
        download_from_url "$PARAKEET_DIR" "${MODEL_BASE_URL}/parakeet-tdt-0.6b-v3.tar.gz"
        download_from_url "$EOU_DIR" "${MODEL_BASE_URL}/realtime_eou_120m-v1-onnx.tar.gz"
        download_from_url "$DIARIZATION_DIR" "${MODEL_BASE_URL}/diarization.tar.gz"
    else
        echo "Downloading models from Hugging Face..."

        # Parakeet TDT from istupakov/parakeet-tdt-0.6b-v3-onnx
        # Required: encoder-model.onnx, encoder-model.onnx.data, decoder_joint-model.onnx, vocab.txt
        download_from_hf "$MODELS_DIR/$PARAKEET_DIR" "istupakov/parakeet-tdt-0.6b-v3-onnx"

        # Verify TDT files were downloaded
        if [ ! -f "$MODELS_DIR/$PARAKEET_DIR/vocab.txt" ]; then
            echo "ERROR: vocab.txt not found in parakeet TDT model"
            echo "Contents of $MODELS_DIR/$PARAKEET_DIR:"
            ls -la "$MODELS_DIR/$PARAKEET_DIR"
            exit 1
        fi

        # EOU model from altunenes/parakeet-rs (subdirectory)
        download_from_hf_subdir "$MODELS_DIR/$EOU_DIR" "altunenes/parakeet-rs" "realtime_eou_120m-v1-onnx"

        # Diarization model from altunenes/parakeet-rs (single file at root)
        mkdir -p "$MODELS_DIR/$DIARIZATION_DIR"
        hf download "altunenes/parakeet-rs" "diar_streaming_sortformer_4spk-v2.1.onnx" --local-dir "$MODELS_DIR/$DIARIZATION_DIR"
    fi

    echo "All models downloaded successfully"
fi

# Download Chatterbox TTS models (for TTS functionality)
CHATTERBOX_MODEL_DIR="${CHATTERBOX_MODEL_DIR:-/app/models/chatterbox-turbo}"

download_chatterbox_tts() {
    if [ -d "$CHATTERBOX_MODEL_DIR" ] && \
       [ -f "$CHATTERBOX_MODEL_DIR/speech_encoder.onnx" ] && \
       [ -f "$CHATTERBOX_MODEL_DIR/embed_tokens.onnx" ] && \
       [ -f "$CHATTERBOX_MODEL_DIR/language_model.onnx" ] && \
       [ -f "$CHATTERBOX_MODEL_DIR/conditional_decoder.onnx" ] && \
       [ -f "$CHATTERBOX_MODEL_DIR/tokenizer.json" ]; then
        echo "Chatterbox TTS models already exist, skipping..."
        return 0
    fi

    echo "Downloading Chatterbox TTS models..."
    mkdir -p "$CHATTERBOX_MODEL_DIR"

    # Download only the required files from ResembleAI/chatterbox-turbo-ONNX
    # (skip fp16, q4, q4f16, quantized variants to save ~4GB)
    echo "Downloading ResembleAI/chatterbox-turbo-ONNX (standard precision only)..."
    local tmpdir=$(mktemp -d)
    hf download ResembleAI/chatterbox-turbo-ONNX \
        --include "onnx/speech_encoder.onnx" \
        --include "onnx/speech_encoder.onnx_data" \
        --include "onnx/embed_tokens.onnx" \
        --include "onnx/embed_tokens.onnx_data" \
        --include "onnx/language_model.onnx" \
        --include "onnx/language_model.onnx_data" \
        --include "onnx/conditional_decoder.onnx" \
        --include "onnx/conditional_decoder.onnx_data" \
        --include "tokenizer.json" \
        --local-dir "$tmpdir"

    # Move files to the model directory (flatten the onnx/ subdirectory)
    mv "$tmpdir"/onnx/*.onnx "$CHATTERBOX_MODEL_DIR"/
    mv "$tmpdir"/onnx/*.onnx_data "$CHATTERBOX_MODEL_DIR"/
    mv "$tmpdir"/tokenizer.json "$CHATTERBOX_MODEL_DIR"/
    rm -rf "$tmpdir"

    echo "Chatterbox TTS models downloaded successfully"
}

download_chatterbox_tts

# Execute the main command
exec "$@"