#!/bin/bash

# Set directories
BASE_DIR="$HOME/voice_training"
PROCESSED_DIR="$BASE_DIR/_02_Processed_Audio"
TRANSCRIPT_DIR="$BASE_DIR/_03_Transcripts"
DATASET_DIR="$BASE_DIR/_04_Training_Dataset"
TRAIN_DIR="$BASE_DIR/_05_Trained_Model"

# Ensure directories exist
mkdir -p "$DATASET_DIR" "$TRAIN_DIR"

# Generate dataset JSON
echo "Generating dataset.json for training..."
python3 - <<EOF
import os
import json

data = []
audio_dir = "$PROCESSED_DIR"
transcript_dir = "$TRANSCRIPT_DIR"
output_json = "$DATASET_DIR/dataset.json"

for file in os.listdir(audio_dir):
    if file.endswith(".wav"):
        filename = os.path.splitext(file)[0]
        transcript_path = os.path.join(transcript_dir, filename + ".txt")

        if os.path.exists(transcript_path):
            with open(transcript_path, "r") as f:
                text = f.read().strip()
            data.append({"audio_filepath": os.path.join(audio_dir, file), "text": text})

with open(output_json, "w") as f:
    json.dump(data, f, indent=4)

print(f"Dataset JSON saved: {output_json}")
EOF

# Start training
echo "Starting Coqui TTS training..."
python3 -m TTS.bin.train_tts --config_path config.json --output_path "$TRAIN_DIR"