#!/bin/bash

# Ensure directories exist
mkdir -p models data/input/reference_audio data/output

# Check if model files exist
if [ ! -f "models/best_model_479919.pth" ] || [ ! -f "models/config.json" ]; then
    echo "ERROR: Model files not found in models/ directory!"
    echo "Please place model files in models/ directory:"
    echo "- models/best_model_479919.pth"
    echo "- models/config.json"
    exit 1
fi

# Check if sample inputs exist
if [ ! -f "data/input/texts.txt" ]; then
    echo "WARNING: No texts.txt found. Creating sample file..."
    echo -e "1\tte\tవడ్రంగి, క్షురక వృత్తులలో పెట్టుబడి ప్రధానమై ఇతరులు కూడా ఈ వృత్తిలో ప్రవేశించి వ్యాపారంగా మార్చేసార\tspeaker1.wav" > data/input/texts.txt
    echo "Please add reference audio files to data/input/reference_audio/"
fi

# Build the Docker image
echo "Building Docker image (this may take some time)..."
docker build -t tts-baseline .

echo ""
echo "Build complete!"
echo ""

docker run --gpus all \
  -v "$(pwd)/models:/app/models" \
  -v "$(pwd)/data/input:/app/data/input" \
  -v "$(pwd)/data/output:/app/data/output" \
  -v "$(pwd)/model_related:/app/model_related" \
  tts-baseline \
  --text_file /app/data/input/texts.txt \
  --ref_dir /app/data/input/reference_audio \
  --savedir /app/data/output \
  --device cuda

# For CPU inference
# docker run \
#   -v "$(pwd)/models:/app/models" \
#   -v "$(pwd)/data/input:/app/data/input" \
#   -v "$(pwd)/data/output:/app/data/output" \
#   -v "$(pwd)/model_related:/app/model_related" \
#   tts-baseline \
#   --text_file /app/data/input/texts.txt \
#   --ref_dir /app/data/input/reference_audio \
#   --savedir /app/data/output \
#   --device cpu