#!/usr/bin/env bash

# Default values
jobs=0  # 0 means auto-detect
output_dir="pdfs"
verbose=0
quiet=0
create_dirs=0
dry_run=0
skip_newer=0
show_progress=1
force=0
select_mode=0
compile_log="typst_compile_errors.log"
move_log="typst_move_errors.log"
declare -a exclude_patterns

# Show usage/help information
show_help() {
    cat << EOF
Usage: $(basename "$0") [OPTIONS]

Compile Typst files and move generated PDFs to a designated directory.

Options:
  -j, --jobs NUM          Number of parallel jobs (default: auto-detect)
  -o, --output-dir DIR    Output directory name (default: pdfs)
  -v, --verbose           Increase verbosity
  -q, --quiet             Suppress most output
  -c, --create-dirs       Create output directories if they don't exist
  -d, --dry-run           Show what would be done without doing it
  -s, --skip-newer        Skip compilation if PDF is newer than source
  -S, --select            Interactive selection mode using skim
  -f, --force             Force compilation even if PDF exists
      --no-progress       Disable progress bar
      --compile-log FILE  Custom location for compilation log (default: $compile_log)
      --move-log FILE     Custom location for move log (default: $move_log)
  -e, --exclude PATTERN   Exclude files matching pattern (can be used multiple times)
  -h, --help              Show this help message and exit

Examples:
  $(basename "$0") -j 4 -o output -c
  $(basename "$0") --verbose --skip-newer
  $(basename "$0") -S                # Select files to compile interactively
  $(basename "$0") -e "**/test/**" -e "**/draft/**"
EOF
}

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -j|--jobs)
            jobs="$2"
            shift 2
            ;;
        -o|--output-dir)
            output_dir="$2"
            shift 2
            ;;
        -v|--verbose)
            verbose=1
            shift
            ;;
        -q|--quiet)
            quiet=1
            shift
            ;;
        -c|--create-dirs)
            create_dirs=1
            shift
            ;;
        -d|--dry-run)
            dry_run=1
            shift
            ;;
        -s|--skip-newer)
            skip_newer=1
            shift
            ;;
        -S|--select)
            select_mode=1
            shift
            ;;
        -f|--force)
            force=1
            shift
            ;;
        --no-progress)
            show_progress=0
            shift
            ;;
        --compile-log)
            compile_log="$2"
            shift 2
            ;;
        --move-log)
            move_log="$2"
            shift 2
            ;;
        -e|--exclude)
            exclude_patterns+=("$2")
            shift 2
            ;;
        -h|--help)
            show_help
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            show_help
            exit 1
            ;;
    esac
done

# Check for conflicting options
if [ "$verbose" -eq 1 ] && [ "$quiet" -eq 1 ]; then
    echo "Error: Cannot use both --verbose and --quiet"
    exit 1
fi

# Check if required tools are installed
check_tool() {
    if ! command -v "$1" &> /dev/null; then
        echo "$1 is not installed. Please install it first."
        echo "On most systems: $2"
        exit 1
    fi
}

check_tool "fd" "cargo install fd-find or apt/brew install fd-find"
check_tool "rg" "cargo install ripgrep or apt/brew install ripgrep"
check_tool "parallel" "apt/brew install parallel"

# Check for skim and bat if select mode is enabled
if [ "$select_mode" -eq 1 ]; then
    check_tool "sk" "cargo install skim or apt/brew install skim"
    check_tool "bat" "cargo install bat or apt/brew install bat"
fi

# ANSI color codes
GREEN='\033[0;32m'
BLUE='\033[0;34m'
YELLOW='\033[0;33m'
RED='\033[0;31m'
CYAN='\033[0;36m'
PURPLE='\033[0;35m'
NC='\033[0m' # No Color
BOLD='\033[1m'

# Nerd font icons
ICON_SUCCESS="󰄬 "
ICON_ERROR="󰅚 "
ICON_WORKING="󰄾 "
ICON_COMPILE="󰈙 "
ICON_MOVE="󰆐 "
ICON_COMPLETE="󰄲 "
ICON_SUMMARY="󰋽 "
ICON_INFO="󰋼 "
ICON_DEBUG="󰃤 "
ICON_SELECT="󰓾 "

# Logging functions
log_debug() {
    if [ "$verbose" -eq 1 ]; then
        echo -e "${BLUE}$ICON_DEBUG${NC} $*"
    fi
}

log_info() {
    if [ "$quiet" -eq 0 ]; then
        echo -e "${CYAN}$ICON_INFO${NC} $*"
    fi
}

log_warning() {
    if [ "$quiet" -eq 0 ]; then
        echo -e "${YELLOW}⚠️ ${NC} $*"
    fi
}

log_error() {
    echo -e "${RED}${BOLD}$ICON_ERROR${NC} $*"
}

log_success() {
    if [ "$quiet" -eq 0 ]; then
        echo -e "${GREEN}${BOLD}$ICON_SUCCESS${NC} $*"
    fi
}

# Create a directory for temporary files
temp_dir=$(mktemp -d)
compile_failures_dir="$temp_dir/compile_failures"
move_failures_dir="$temp_dir/move_failures"
progress_dir="$temp_dir/progress"
mkdir -p "$compile_failures_dir" "$move_failures_dir" "$progress_dir"

# Create a lock file for progress updates and a flag for final progress
progress_lock="/tmp/typst_progress_lock"
final_progress_file="$temp_dir/final_progress"
touch "$final_progress_file"

# Initialize log files
> "$compile_log"
> "$move_log"

# Store current working directory
CWD=$(pwd)

log_info "Starting Typst compilation process..."

# Build exclude arguments for fd
fd_exclude_args=()
for pattern in "${exclude_patterns[@]}"; do
    fd_exclude_args+=("-E" "$pattern")
done

# Create a list of files to process
typst_files_list="$temp_dir/typst_files.txt"

# Create a custom bat configuration for Typst syntax
setup_bat_for_typst() {
    bat_config_dir="$temp_dir/bat_config"
    mkdir -p "$bat_config_dir/syntaxes"

    # Create a basic syntax mapping file for Typst
    cat > "$bat_config_dir/syntaxes/typst.sublime-syntax" << 'TYPST_SYNTAX'
%YAML 1.2
---
name: Typst
file_extensions:
  - typ
scope: source.typst

contexts:
  main:
    # Comments
    - match: /\*
      scope: comment.block.typst
      push: block_comment

    - match: //
      scope: comment.line.double-slash.typst
      push: line_comment

    # Strings
    - match: '"'
      scope: punctuation.definition.string.begin.typst
      push: double_string

    # Math
    - match: '\$'
      scope: punctuation.definition.math.begin.typst
      push: math

    # Functions
    - match: '#([a-zA-Z][a-zA-Z0-9_]*)'
      scope: entity.name.function.typst

    # Variables
    - match: '\b([a-zA-Z][a-zA-Z0-9_]*)\s*:'
      scope: variable.other.typst

    # Keywords
    - match: '\b(let|set|show|if|else|for|in|while|return|import|include|at|do|not|and|or|none|auto)\b'
      scope: keyword.control.typst

  block_comment:
    - match: \*/
      scope: comment.block.typst
      pop: true
    - match: .
      scope: comment.block.typst

  line_comment:
    - match: $
      pop: true
    - match: .
      scope: comment.line.double-slash.typst

  double_string:
    - match: '"'
      scope: punctuation.definition.string.end.typst
      pop: true
    - match: \\.
      scope: constant.character.escape.typst
    - match: .
      scope: string.quoted.double.typst

  math:
    - match: '\$'
      scope: punctuation.definition.math.end.typst
      pop: true
    - match: .
      scope: markup.math.typst
TYPST_SYNTAX

    echo "$bat_config_dir"
}

if [ "$select_mode" -eq 1 ]; then
    log_info "${PURPLE}$ICON_SELECT${NC} Interactive selection mode enabled"
    log_info "Use TAB to select multiple files, ENTER to confirm"

    # Set up bat configuration for Typst syntax highlighting
    bat_config=$(setup_bat_for_typst)
    export BAT_CONFIG_PATH="$bat_config"

    # Use skim with bat for preview to select files
    selected_files="$temp_dir/selected_files.txt"

    # Find all eligible .typ files for selection
    fd '\.typ$' --type f "${fd_exclude_args[@]}" . > "$typst_files_list.all"

    # Check if we found any files
    if [ ! -s "$typst_files_list.all" ]; then
        log_error "No .typ files found matching your criteria."
        rm -rf "$temp_dir"
        exit 0
    fi

    # Prepare preview command for skim: use bat with custom syntax for .typ files
    preview_cmd="bat --color=always --style=numbers --map-syntax='*.typ:Markdown' {} 2>/dev/null || cat {}"

    # Use skim with bat preview to select files
    cat "$typst_files_list.all" | sk --multi \
        --preview "$preview_cmd" \
        --preview-window "right:70%" \
        --height "80%" \
        --prompt "Select .typ files to compile: " \
        --header "TAB: Select multiple files, ENTER: Confirm, CTRL-C: Cancel" \
        --no-mouse > "$selected_files"

    # Check if user selected any files
    if [ ! -s "$selected_files" ]; then
        log_error "No files selected. Exiting."
        rm -rf "$temp_dir"
        exit 0
    fi

    # Use the selected files instead of all discovered files
    cp "$selected_files" "$typst_files_list"
    total_selected=$(wc -l < "$typst_files_list")
    total_available=$(wc -l < "$typst_files_list.all")

    log_info "Selected ${BOLD}$total_selected${NC} out of ${BOLD}$total_available${NC} available files"

    # Display tip about better Typst syntax highlighting
    cat << 'TYPST_TIP'

📝 To get proper Typst syntax highlighting in bat:

1. Create a custom Typst syntax file:
   mkdir -p ~/.config/bat/syntaxes
   curl -L https://raw.githubusercontent.com/typst/typst-vs-code/main/syntaxes/typst.tmLanguage.json \
     -o ~/.config/bat/syntaxes/typst.tmLanguage.json

2. Build bat's syntax cache:
   bat cache --build

This will provide proper syntax highlighting for Typst files in future uses!

TYPST_TIP
else
    # Normal mode - process all files
    fd '\.typ$' --type f "${fd_exclude_args[@]}" . > "$typst_files_list"
    log_info "Found ${BOLD}$(wc -l < "$typst_files_list")${NC} Typst files to process"
fi

total_files=$(wc -l < "$typst_files_list")

# Function to update progress bar during processing
update_progress_during() {
    # If progress is disabled, do nothing
    if [ "$show_progress" -eq 0 ]; then
        return 0
    fi

    (
        # Try to acquire lock, but don't wait if busy
        flock -n 200 || return 0

        completed=$(find "$progress_dir" -type f | wc -l)
        percent=$((completed * 100 / total_files))
        bar_length=50
        filled_length=$((bar_length * completed / total_files))

        # Create the progress bar
        bar=""
        for ((i=0; i<bar_length; i++)); do
            if [ $i -lt $filled_length ]; then
                bar="${bar}█"
            else
                bar="${bar}░"
            fi
        done

        # Calculate success and failure counts
        success=$((completed - $(find "$compile_failures_dir" "$move_failures_dir" -type f | wc -l)))
        compile_fails=$(find "$compile_failures_dir" -type f | wc -l)
        move_fails=$(find "$move_failures_dir" -type f | wc -l)

        # Clear the previous line and print the updated progress
        echo -ne "\r\033[K"
        echo -ne "${PURPLE}$ICON_WORKING Progress: ${GREEN}$bar ${BOLD}${percent}%${NC} "
        echo -ne "[${GREEN}${success}✓${NC}|${RED}${compile_fails}✗${NC}|${YELLOW}${move_fails}!${NC}] "
        echo -ne "${BLUE}($completed/$total_files)${NC}"

        # If all files are processed, print the completion message ONLY ONCE
        if [ $completed -eq $total_files ] && [ ! -e "$final_progress_file.done" ]; then
            touch "$final_progress_file.done"
            echo -e "\n${GREEN}${BOLD}$ICON_COMPLETE All files processed!${NC}"
        fi
    ) 200>"$progress_lock"
}

# Function to show final progress (called only once at the end)
show_final_progress() {
    # If progress is disabled, do nothing
    if [ "$show_progress" -eq 0 ]; then
        return 0
    fi

    (
        flock -w 1 200

        # Only proceed if the final progress hasn't been shown yet
        if [ -e "$final_progress_file.done" ]; then
            return 0
        fi

        completed=$(find "$progress_dir" -type f | wc -l)
        percent=$((completed * 100 / total_files))
        bar_length=50
        filled_length=$((bar_length * completed / total_files))

        # Create the progress bar
        bar=""
        for ((i=0; i<bar_length; i++)); do
            if [ $i -lt $filled_length ]; then
                bar="${bar}█"
            else
                bar="${bar}░"
            fi
        done

        # Calculate success and failure counts
        success=$((completed - $(find "$compile_failures_dir" "$move_failures_dir" -type f | wc -l)))
        compile_fails=$(find "$compile_failures_dir" -type f | wc -l)
        move_fails=$(find "$move_failures_dir" -type f | wc -l)

        # Clear the previous line and print the updated progress
        echo -ne "\r\033[K"
        echo -ne "${PURPLE}$ICON_WORKING Progress: ${GREEN}$bar ${BOLD}${percent}%${NC} "
        echo -ne "[${GREEN}${success}✓${NC}|${RED}${compile_fails}✗${NC}|${YELLOW}${move_fails}!${NC}] "
        echo -ne "${BLUE}($completed/$total_files)${NC}"

        # Mark final progress as shown
        touch "$final_progress_file.done"

        # If all files are processed, print the completion message
        if [ $completed -eq $total_files ]; then
            echo -e "\n${GREEN}${BOLD}$ICON_COMPLETE All files processed!${NC}"
        fi
    ) 200>"$progress_lock"
}

# Function to process a single .typ file
process_file() {
    typfile="$1"
    file_id=$(echo "$typfile" | md5sum | cut -d' ' -f1)

    # Get the directory containing the .typ file
    typdir=$(dirname "$typfile")
    # Get the filename without path
    filename=$(basename "$typfile")
    # Get the filename without extension
    basename="${filename%.typ}"

    # Check if output directory exists or should be created
    target_dir="$typdir/$output_dir"
    if [ ! -d "$target_dir" ]; then
        if [ "$create_dirs" -eq 1 ]; then
            if [ "$dry_run" -eq 0 ]; then
                mkdir -p "$target_dir"
                log_debug "Created directory: $target_dir"
            else
                log_debug "[DRY RUN] Would create directory: $target_dir"
            fi
        else
            # Skip this file if output directory doesn't exist and --create-dirs not specified
            log_debug "Skipping $typfile (no $output_dir directory)"
            touch "$progress_dir/$file_id"
            update_progress_during
            return 0
        fi
    fi

    # Skip if PDF is newer than source and --skip-newer is specified
    if [ "$skip_newer" -eq 1 ] && [ -f "$target_dir/$basename.pdf" ]; then
        if [ "$typfile" -ot "$target_dir/$basename.pdf" ] && [ "$force" -eq 0 ]; then
            log_debug "Skipping $typfile (PDF is newer)"
            touch "$progress_dir/$file_id"
            update_progress_during
            return 0
        fi
    fi

    # Create a temporary file for capturing compiler output
    temp_output="$temp_dir/output_${file_id}.log"

    # Add a header to the log before compilation
    {
        echo -e "\n===== COMPILING: $typfile ====="
        echo "$(date)"
    } > "$temp_output.header"

    # In dry run mode, just log what would be done
    if [ "$dry_run" -eq 1 ]; then
        log_debug "[DRY RUN] Would compile: $typfile"
        log_debug "[DRY RUN] Would move to: $target_dir/$basename.pdf"
        touch "$progress_dir/$file_id"
        update_progress_during
        return 0
    fi

    # Compile the .typ file using typst with --root flag and capture all output
    if ! typst compile --root "$CWD" "$typfile" > "$temp_output.stdout" 2> "$temp_output.stderr"; then
        # Store the failure
        echo "$typfile" > "$compile_failures_dir/$file_id"
        log_debug "Compilation failed for $typfile"

        # Combine stdout and stderr
        cat "$temp_output.stdout" "$temp_output.stderr" > "$temp_output.combined"

        # Filter the output to only include error messages using ripgrep
        rg "error:" -A 20 "$temp_output.combined" > "$temp_output.errors" || true

        # Lock the log file to avoid concurrent writes corrupting it
        (
            flock -w 1 201
            cat "$temp_output.header" "$temp_output.errors" >> "$compile_log"
            echo -e "\n" >> "$compile_log"
        ) 201>"$compile_log.lock"
    else
        # Check if the output PDF exists
        if [ -f "$typdir/$basename.pdf" ]; then
            # Try to move the output PDF to the output directory
            move_header="$temp_dir/move_${file_id}.header"
            {
                echo -e "\n===== MOVING: $typfile ====="
                echo "$(date)"
            } > "$move_header"

            if ! mv "$typdir/$basename.pdf" "$target_dir/" 2> "$temp_output.move_err"; then
                echo "$typfile -> $target_dir/$basename.pdf" > "$move_failures_dir/$file_id"
                log_debug "Failed to move $basename.pdf to $target_dir/"

                # Lock the log file to avoid concurrent writes corrupting it
                (
                    flock -w 1 202
                    cat "$move_header" "$temp_output.move_err" >> "$move_log"
                    echo "Failed to move $typdir/$basename.pdf to $target_dir/" >> "$move_log"
                ) 202>"$move_log.lock"
            else
                log_debug "Moved $basename.pdf to $target_dir/"
            fi
        else
            # This is a fallback check in case typst doesn't return error code
            echo "$typfile" > "$compile_failures_dir/$file_id"
            log_debug "Compilation completed without errors but no PDF was generated for $typfile"

            # Lock the log file to avoid concurrent writes corrupting it
            (
                flock -w 1 201
                echo "Compilation completed without errors but no PDF was generated" >> "$compile_log"
            ) 201>"$compile_log.lock"
        fi
    fi

    # Mark this file as processed (for progress tracking)
    touch "$progress_dir/$file_id"

    # Update the progress bar
    update_progress_during
}

export -f process_file
export -f update_progress_during
export -f show_final_progress
export -f log_debug
export -f log_info
export -f log_warning
export -f log_error
export -f log_success
export CWD
export temp_dir
export compile_failures_dir
export move_failures_dir
export progress_dir
export final_progress_file
export progress_lock
export compile_log
export move_log
export total_files
export GREEN BLUE YELLOW RED CYAN PURPLE NC BOLD
export ICON_SUCCESS ICON_ERROR ICON_WORKING ICON_COMPILE ICON_MOVE ICON_COMPLETE ICON_SUMMARY ICON_INFO ICON_DEBUG
export verbose
export quiet
export output_dir
export create_dirs
export dry_run
export skip_newer
export show_progress
export force

# Determine the number of CPU cores and use that many parallel jobs (if not specified)
if [ "$jobs" -eq 0 ]; then
    jobs=$(nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 4)
fi
log_info "Using ${BOLD}$jobs${NC} parallel jobs for compilation"

# Initialize progress bar if showing progress
if [ "$show_progress" -eq 1 ]; then
    update_progress_during
fi

# Process files in parallel with --will-cite to suppress citation notice
if [ "$total_files" -gt 0 ]; then
    cat "$typst_files_list" | parallel --will-cite --jobs "$jobs" process_file

    # Wait a moment for any remaining progress updates to complete
    sleep 0.5

    # Show the final progress exactly once if showing progress
    if [ "$show_progress" -eq 1 ]; then
        show_final_progress
    fi

    # Print summary of failures
    if [ "$quiet" -eq 0 ]; then
        echo -e "\n${BOLD}${PURPLE}$ICON_SUMMARY Processing Summary${NC}"
    fi

    # Collect all failure files
    compile_failures=$(find "$compile_failures_dir" -type f | wc -l)
    move_failures=$(find "$move_failures_dir" -type f | wc -l)

    if [ "$compile_failures" -eq 0 ] && [ "$move_failures" -eq 0 ]; then
        log_success "All files processed successfully."
    else
        if [ "$compile_failures" -gt 0 ]; then
            echo -e "\n${RED}${BOLD}$ICON_ERROR Compilation failures:${NC}"
            find "$compile_failures_dir" -type f -exec cat {} \; | sort | while read -r failure; do
                echo -e "${RED}- $failure${NC}"
            done
            echo -e "${BLUE}See $compile_log for detailed error messages.${NC}"
        fi

        if [ "$move_failures" -gt 0 ]; then
            echo -e "\n${YELLOW}${BOLD}$ICON_ERROR Move failures:${NC}"
            find "$move_failures_dir" -type f -exec cat {} \; | sort | while read -r failure; do
                echo -e "${YELLOW}- $failure${NC}"
            done
            echo -e "${BLUE}See $move_log for detailed error messages.${NC}"
        fi
    fi
else
    log_warning "No .typ files found to process."
fi

# Clean up temporary directory and lock files
rm -rf "$temp_dir"
rm -f "$progress_lock" "$compile_log.lock" "$move_log.lock"

log_success "Processing complete."