#!/bin/bash
set -euo pipefail

if [[ $# -ne 2 ]]; then
    echo "Usage: $0 <source_folder> <target_folder>" >&2
    exit 1
fi
SOURCE_DIR=$(realpath "$1")  # Normalize source path
readonly TARGET_DIR="$2"
readonly TEMP_ERROR=$(mktemp)
readonly LOG_FILE="/tmp/convert2md_$$.log"
readonly UNPROCESSED_FILE="/tmp/unprocessed_$$.txt"
trap 'rm -f "$TEMP_ERROR" "$LOG_FILE" "$UNPROCESSED_FILE"' EXIT

validate_inputs() {
    [[ ! -d "$SOURCE_DIR" ]] && { echo "Error: Source directory '$SOURCE_DIR' does not exist." >&2; exit 1; }
    mkdir -p "$TARGET_DIR" || { echo "Error: Could not create target directory '$TARGET_DIR'." >&2; exit 1; }
    command -v rst2myst >/dev/null 2>&1 || { echo "Error: rst2myst not installed. Use 'pip install rst-to-myst'." >&2; exit 1; }
}

process_rst_file() {
    local rst_file="$1"
    local relative_path="${rst_file#$SOURCE_DIR/}"
    local md_file_name="${relative_path%.rst}.md"
    local target_file="$TARGET_DIR/$md_file_name"
    local target_dir=$(dirname "$target_file")
    mkdir -p "$target_dir" || { echo "Error: Could not create '$target_dir' for '$rst_file'" >&2; echo "$rst_file" >> "$UNPROCESSED_FILE"; return 1; }
    if rst2myst stream --config config.yaml "$rst_file" > "$target_file" 2>>"$TEMP_ERROR"; then
        # Replace 'rst-class' with 'container' and '.rst' with '.md' in content
        sed -i 's/rst-class/container/g; s/\.rst/.md/g' "$target_file" 2>>"$TEMP_ERROR" || {
            echo "Error: Failed to modify content of '$target_file'" >&2
            echo "$rst_file" >> "$UNPROCESSED_FILE"
            return 1
        }
        echo "Converted: $rst_file -> $target_file" | tee -a "$LOG_FILE"
        return 0
    else
        echo "Failed to convert: $rst_file" >&2
        cat "$TEMP_ERROR" >&2
        echo "$rst_file" >> "$UNPROCESSED_FILE"
        return 1
    fi
}

copy_non_rst_file() {
    local file="$1"
    local relative_path="${file#$SOURCE_DIR/}"
    local target_file="$TARGET_DIR/$relative_path"
    local target_dir=$(dirname "$target_file")
    mkdir -p "$target_dir" && cp -p "$file" "$target_file" 2>>"$TEMP_ERROR" && {
        # Replace 'rst-class' with 'container' and '.rst' with '.md' in content
        sed -i 's/rst-class/container/g; s/\.rst/.md/g' "$target_file" 2>>"$TEMP_ERROR" || {
            echo "Error: Failed to modify content of '$target_file'" >&2
            echo "$file" >> "$UNPROCESSED_FILE"
            return 1
        }
        echo "Copied: $file -> $target_file" | tee -a "$LOG_FILE"
        return 0
    } || {
        echo "Failed to copy: $file" >&2
        echo "$file" >> "$UNPROCESSED_FILE"
        return 1
    }
}

main() {
    validate_inputs
    echo "Starting conversion process..." | tee "$LOG_FILE"
    
    # Log all files found
    mapfile -t RST_FILES < <(find "$SOURCE_DIR" -type f -name "*.rst" 2>>"$LOG_FILE")
    mapfile -t NON_RST_FILES < <(find "$SOURCE_DIR" -type f ! -name "*.rst" 2>>"$LOG_FILE")
    echo "Found ${#RST_FILES[@]} RST files and ${#NON_RST_FILES[@]} non-RST files" | tee -a "$LOG_FILE"
    printf '%s\n' "${RST_FILES[@]}" | sed 's/^/RST: /' >> "$LOG_FILE"
    printf '%s\n' "${NON_RST_FILES[@]}" | sed 's/^/Non-RST: /' >> "$LOG_FILE"
    
    # Initialize unprocessed file list with all found files
    printf '%s\n' "${RST_FILES[@]}" "${NON_RST_FILES[@]}" > "$UNPROCESSED_FILE"
    
    local parallel_jobs=$(( $(nproc) / 2 ))
    export -f process_rst_file copy_non_rst_file
    export SOURCE_DIR TARGET_DIR TEMP_ERROR LOG_FILE UNPROCESSED_FILE
    
    # Process RST files
    printf '%s\0' "${RST_FILES[@]}" | xargs -0 -P "$parallel_jobs" -I {} bash -c 'process_rst_file "{}"' || {
        echo "Warning: Some RST processing failed. Check $LOG_FILE." >&2
    }
    
    # Process non-RST files
    printf '%s\0' "${NON_RST_FILES[@]}" | xargs -0 -P "$parallel_jobs" -I {} bash -c 'copy_non_rst_file "{}"' || {
        echo "Warning: Some file copies failed. Check $LOG_FILE." >&2
    }
    
    # Summary
    local successful=$(grep -c "^Converted:" "$LOG_FILE")
    local failed=$(grep -c "^Failed to convert:" "$LOG_FILE")
    local copied=$(grep -c "^Copied:" "$LOG_FILE")
    echo "Conversion complete." | tee -a "$LOG_FILE"
    echo "RST files found: ${#RST_FILES[@]}" | tee -a "$LOG_FILE"
    echo "RST files converted: $successful" | tee -a "$LOG_FILE"
    echo "RST files failed: $failed" | tee -a "$LOG_FILE"
    echo "Non-RST files copied: $copied" | tee -a "$LOG_FILE"
    
    # List unprocessed files (those still in UNPROCESSED_FILE after processing)
    if [[ -s "$UNPROCESSED_FILE" ]]; then
        echo "Files not processed:" | tee -a "$LOG_FILE"
        cat "$UNPROCESSED_FILE" | while IFS= read -r file; do
            echo "  $file" | tee -a "$LOG_FILE"
        done
        echo "Note: Some files were not processed. See above list and $LOG_FILE for details." >&2
        exit 1
    else
        echo "All files processed successfully." | tee -a "$LOG_FILE"
    fi
    
    if [[ ${#RST_FILES[@]} -ne $((successful + failed)) ]]; then
        echo "Error: ${#RST_FILES[@]} RST files found, but only $((successful + failed)) processed!" >&2
        exit 1
    fi
    [[ $failed -gt 0 ]] && { echo "Note: $failed RST conversions failed. See $LOG_FILE." >&2; exit 1; }
}

main "$@"