#!/usr/bin/env bash
# tree.sh - Directory tree generator
# Generates directory tree structure with optional file listing,
# markdown export, and statistics.
#
# Usage: ./tree.sh [-p PATH] [-o OUTPUT] [-d DEPTH] [-f] [-s] [-h] [-v]
#
# Version: 1.0.0

set -euo pipefail

# ── Constants ────────────────────────────────────────────────────────────────
readonly VERSION="1.0.0"
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME
readonly DEFAULT_OUTPUT="tree_output.md"
readonly IGNORE_FILE=".treeignore"

# Error codes
readonly ERR_OK=0
readonly ERR_USAGE=1
readonly ERR_PATH=2
readonly ERR_WRITE=3
# shellcheck disable=SC2034
readonly ERR_INTERNAL=4  # Fallback for unhandled errors

# ── Global State ─────────────────────────────────────────────────────────────
TARGET_PATH=""
OUTPUT_FILE="${DEFAULT_OUTPUT}"
MAX_DEPTH=0
SHOW_FILES=false
SHOW_STATS=true
declare -a IGNORE_PATTERNS=()
declare -A SEEN_INODES=()

# Counters for statistics
DIR_COUNT=0
FILE_COUNT=0
TOTAL_SIZE=0

# ── Module: parse_args ──────────────────────────────────────────────────────
# Parse command-line arguments and populate global state.
# Returns 0 on success, 1 on invalid arguments.
parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            -p|--path)
                if [[ $# -lt 2 ]]; then
                    echo "Error: -p/--path requires an argument" >&2
                    return "$ERR_USAGE"
                fi
                TARGET_PATH="$2"
                shift 2
                ;;
            -o|--output)
                if [[ $# -lt 2 ]]; then
                    echo "Error: -o/--output requires an argument" >&2
                    return "$ERR_USAGE"
                fi
                OUTPUT_FILE="$2"
                shift 2
                ;;
            -d|--depth)
                if [[ $# -lt 2 ]]; then
                    echo "Error: -d/--depth requires an argument" >&2
                    return "$ERR_USAGE"
                fi
                if ! [[ "$2" =~ ^[0-9]+$ ]] || [[ "$2" -eq 0 ]]; then
                    echo "Error: depth must be a positive integer" >&2
                    return "$ERR_USAGE"
                fi
                MAX_DEPTH="$2"
                shift 2
                ;;
            -f|--files)
                SHOW_FILES=true
                shift
                ;;
            -s|--no-stats)
                SHOW_STATS=false
                shift
                ;;
            -h|--help)
                show_help
                exit "$ERR_OK"
                ;;
            -v|--version)
                echo "${SCRIPT_NAME} ${VERSION}"
                exit "$ERR_OK"
                ;;
            -*)
                echo "Error: unknown option '$1'" >&2
                echo "Run '${SCRIPT_NAME} --help' for usage." >&2
                return "$ERR_USAGE"
                ;;
            *)
                # Positional argument: treat as path
                if [[ -z "$TARGET_PATH" ]]; then
                    TARGET_PATH="$1"
                else
                    echo "Error: unexpected argument '$1'" >&2
                    return "$ERR_USAGE"
                fi
                shift
                ;;
        esac
    done

    # Default to current directory if no path specified
    if [[ -z "$TARGET_PATH" ]]; then
        TARGET_PATH="."
    fi

    return "$ERR_OK"
}

# ── Module: show_help ───────────────────────────────────────────────────────
show_help() {
    cat <<EOF
${SCRIPT_NAME} ${VERSION} - Directory tree generator

Usage: ${SCRIPT_NAME} [OPTIONS] [PATH]

Options:
  -p, --path <PATH>       Target directory (default: current directory)
  -o, --output <FILE>     Output file for markdown (default: ${DEFAULT_OUTPUT})
  -d, --depth <N>         Maximum recursion depth (default: unlimited)
  -f, --files             Also generate file tree
  -s, --no-stats          Suppress statistics output
  -h, --help              Show this help message
  -v, --version           Show version information

Examples:
  ${SCRIPT_NAME}                      # Tree of current directory
  ${SCRIPT_NAME} -p /home/user        # Tree of /home/user
  ${SCRIPT_NAME} -p . -f -d 3        # Tree with files, depth 3
  ${SCRIPT_NAME} -o mytree.md         # Save to mytree.md
EOF
}

# ── Module: load_ignore_config ──────────────────────────────────────────────
# Load ignore patterns from built-in defaults and .treeignore file.
# Populates the global IGNORE_PATTERNS array.
load_ignore_config() {
    # Built-in default ignore patterns
    IGNORE_PATTERNS=(
        ".git"
        ".svn"
        ".hg"
        "node_modules"
        "__pycache__"
        ".DS_Store"
        "*.pyc"
        "*.pyo"
        ".cache"
        ".tox"
        ".eggs"
        "*.egg-info"
        "dist"
        "build"
        ".next"
        ".nuxt"
        ".output"
        ".vercel"
        ".terraform"
        ".vagrant"
        ".idea"
        ".vscode"
        "*.swp"
        "*.swo"
        "*.swn"
        "*.class"
        "*.o"
        "*.so"
        "*.dylib"
    )

    # Load patterns from .treeignore in the target directory
    local ignore_file
    ignore_file="$(resolve_path "${TARGET_PATH}")/${IGNORE_FILE}"
    if [[ -f "$ignore_file" ]]; then
        local line
        while IFS= read -r line || [[ -n "$line" ]]; do
            # Skip empty lines and comments
            line="$(echo "$line" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')"
            if [[ -n "$line" && "$line" != \#* ]]; then
                IGNORE_PATTERNS+=("$line")
            fi
        done < "$ignore_file"
    fi
}

# ── Module: resolve_path ────────────────────────────────────────────────────
# Resolve a path to its absolute form.
# Arguments: $1 = path to resolve
# Outputs: resolved absolute path
resolve_path() {
    local path="$1"
    if [[ "$path" = /* ]]; then
        echo "$path"
    else
        echo "$(pwd)/${path}"
    fi
}

# ── Module: should_ignore ───────────────────────────────────────────────────
# Check if a given name matches any ignore pattern.
# Arguments: $1 = file or directory name
# Returns: 0 if should be ignored, 1 otherwise
should_ignore() {
    local name="$1"
    local pattern
    for pattern in "${IGNORE_PATTERNS[@]}"; do
        # Direct match
        if [[ "$name" == "$pattern" ]]; then
            return 0
        fi
        # Glob pattern match (e.g., *.pyc)
        # shellcheck disable=SC2254
        case "$name" in
            $pattern) return 0 ;;
        esac
    done
    return 1
}

# ── Module: traverse_directory ──────────────────────────────────────────────
# Recursively traverse directory and build tree structure.
# Arguments:
#   $1 = directory path
#   $2 = current depth
#   $3 = prefix string for tree lines
# Outputs: tree lines to stdout
traverse_directory() {
    local dir="$1"
    local depth="$2"
    local prefix="$3"

    # Check depth limit
    if [[ $MAX_DEPTH -gt 0 && $depth -ge $MAX_DEPTH ]]; then
        return 0
    fi

    # Check if directory is readable
    if [[ ! -r "$dir" ]]; then
        echo "${prefix}[Permission denied]" >&2
        return 0
    fi

    # Collect entries, sorted
    local entries=()
    local entry
    for entry in "$dir"/* "$dir"/.*; do
        # Skip . and ..
        local basename
        basename="$(basename "$entry")"
        [[ "$basename" == "." || "$basename" == ".." ]] && continue
        # Skip if doesn't exist (glob didn't match)
        [[ ! -e "$entry" ]] && continue
        # Check ignore
        if should_ignore "$basename"; then
            continue
        fi
        entries+=("$entry")
    done

    # Sort entries: directories first, then files, both alphabetical
    local dirs=()
    local files=()
    local e
    for e in "${entries[@]+"${entries[@]}"}"; do
        if [[ -d "$e" ]]; then
            dirs+=("$e")
        else
            files+=("$e")
        fi
    done

    # Combine: dirs first, then files
    local sorted=()
    sorted+=("${dirs[@]+"${dirs[@]}"}")
    sorted+=("${files[@]+"${files[@]}"}")

    local total=${#sorted[@]}
    local idx=0

    for e in "${sorted[@]+"${sorted[@]}"}"; do
        idx=$((idx + 1))
        local basename
        basename="$(basename "$e")"
        local is_last=false
        if [[ $idx -eq $total ]]; then
            is_last=true
        fi

        local connector
        local child_prefix
        if $is_last; then
            connector="└── "
            child_prefix="${prefix}    "
        else
            connector="├── "
            child_prefix="${prefix}│   "
        fi

        if [[ -d "$e" ]]; then
            # Inode-based cycle detection
            local inode
            inode="$(stat -c '%d:%i' "$e" 2>/dev/null || stat -f '%d:%i' "$e" 2>/dev/null || echo "")"
            if [[ -n "$inode" && -n "${SEEN_INODES[$inode]:-}" ]]; then
                echo "${prefix}${connector}${basename}/ [cycle detected, skipped]"
                continue
            fi
            if [[ -n "$inode" ]]; then
                SEEN_INODES["$inode"]=1
            fi

            echo "${prefix}${connector}${basename}/"
            DIR_COUNT=$((DIR_COUNT + 1))
            traverse_directory "$e" $((depth + 1)) "$child_prefix"
        else
            echo "${prefix}${connector}${basename}"
            FILE_COUNT=$((FILE_COUNT + 1))
            # Accumulate file size
            local fsize
            fsize="$(stat -c '%s' "$e" 2>/dev/null || stat -f '%z' "$e" 2>/dev/null || echo 0)"
            TOTAL_SIZE=$((TOTAL_SIZE + fsize))
        fi
    done
}

# ── Module: render_tree ─────────────────────────────────────────────────────
# Render the directory tree starting from the target path.
# Outputs: complete tree to stdout
render_tree() {
    local abs_path
    abs_path="$(resolve_path "$TARGET_PATH")"
    local root_name
    root_name="$(basename "$abs_path")"

    # Reset inode tracking and counters for each render pass
    SEEN_INODES=()
    DIR_COUNT=0
    FILE_COUNT=0
    TOTAL_SIZE=0

    echo "${root_name}/"
    DIR_COUNT=$((DIR_COUNT + 1))
    # Track root inode for cycle detection
    local root_inode
    root_inode="$(stat -c '%d:%i' "$abs_path" 2>/dev/null || stat -f '%d:%i' "$abs_path" 2>/dev/null || echo "")"
    if [[ -n "$root_inode" ]]; then
        SEEN_INODES["$root_inode"]=1
    fi

    traverse_directory "$abs_path" 1 ""
}

# ── Module: collect_files ───────────────────────────────────────────────────
# Collect all files recursively from the target directory.
# Arguments:
#   $1 = directory path
#   $2 = current depth
# Outputs: file paths to stdout
collect_files() {
    local dir="$1"
    local depth="$2"

    if [[ $MAX_DEPTH -gt 0 && $depth -ge $MAX_DEPTH ]]; then
        return 0
    fi

    if [[ ! -r "$dir" ]]; then
        return 0
    fi

    local entry
    for entry in "$dir"/* "$dir"/.*; do
        local basename
        basename="$(basename "$entry")"
        [[ "$basename" == "." || "$basename" == ".." ]] && continue
        [[ ! -e "$entry" ]] && continue
        if should_ignore "$basename"; then
            continue
        fi

        if [[ -d "$entry" ]]; then
            collect_files "$entry" $((depth + 1))
        else
            echo "$entry"
        fi
    done
}

# ── Module: render_file_list ────────────────────────────────────────────────
# Render the complete file list.
# Outputs: file list to stdout
render_file_list() {
    local abs_path
    abs_path="$(resolve_path "$TARGET_PATH")"

    echo "Files:"
    local file
    while IFS= read -r file; do
        echo "  ${file}"
    done < <(collect_files "$abs_path" 0 | sort)
}

# ── Module: compute_stats ───────────────────────────────────────────────────
# Compute and output statistics.
# Outputs: statistics to stdout
compute_stats() {
    local size_human
    if [[ $TOTAL_SIZE -ge 1073741824 ]]; then
        size_human="$(awk "BEGIN {printf \"%.2f GB\", ${TOTAL_SIZE}/1073741824}")"
    elif [[ $TOTAL_SIZE -ge 1048576 ]]; then
        size_human="$(awk "BEGIN {printf \"%.2f MB\", ${TOTAL_SIZE}/1048576}")"
    elif [[ $TOTAL_SIZE -ge 1024 ]]; then
        size_human="$(awk "BEGIN {printf \"%.2f KB\", ${TOTAL_SIZE}/1024}")"
    else
        size_human="${TOTAL_SIZE} B"
    fi

    echo ""
    echo "Statistics:"
    echo "  Directories: ${DIR_COUNT}"
    echo "  Files:       ${FILE_COUNT}"
    echo "  Total size:  ${size_human}"
}

# ── Module: output_terminal ─────────────────────────────────────────────────
# Output tree to terminal (stdout).
output_terminal() {
    render_tree
    if $SHOW_FILES; then
        echo ""
        render_file_list
    fi
    if $SHOW_STATS; then
        compute_stats
    fi
}

# ── Module: save_markdown ───────────────────────────────────────────────────
# Save tree output to a Markdown file.
# Arguments: $1 = output file path
save_markdown() {
    local outfile="$1"
    local outfile_dir
    outfile_dir="$(dirname "$outfile")"

    # Ensure output directory exists
    if [[ ! -d "$outfile_dir" ]]; then
        mkdir -p "$outfile_dir" 2>/dev/null || {
            echo "Error: cannot create output directory '${outfile_dir}'" >&2
            exit "$ERR_WRITE"
        }
    fi

    {
        echo "# Directory Tree"
        echo ""
        echo "Path: \`${TARGET_PATH}\`"
        echo ""
        echo '```'
        render_tree
        echo '```'
        if $SHOW_FILES; then
            echo ""
            echo "## Files"
            echo ""
            render_file_list
        fi
        if $SHOW_STATS; then
            echo ""
            echo "## Statistics"
            echo ""
            echo "- **Directories:** ${DIR_COUNT}"
            echo "- **Files:** ${FILE_COUNT}"
            local size_human
            if [[ $TOTAL_SIZE -ge 1073741824 ]]; then
                size_human="$(awk "BEGIN {printf \"%.2f GB\", ${TOTAL_SIZE}/1073741824}")"
            elif [[ $TOTAL_SIZE -ge 1048576 ]]; then
                size_human="$(awk "BEGIN {printf \"%.2f MB\", ${TOTAL_SIZE}/1048576}")"
            elif [[ $TOTAL_SIZE -ge 1024 ]]; then
                size_human="$(awk "BEGIN {printf \"%.2f KB\", ${TOTAL_SIZE}/1024}")"
            else
                size_human="${TOTAL_SIZE} B"
            fi
            echo "- **Total size:** ${size_human}"
        fi
    } > "$outfile" 2>/dev/null || {
        echo "Error: cannot write to '${outfile}'" >&2
        exit "$ERR_WRITE"
    }

    echo "Markdown saved to: ${outfile}"
}

# ── Main ─────────────────────────────────────────────────────────────────────
main() {
    # Parse arguments
    local parse_rc=0
    parse_args "$@" || parse_rc=$?
    if [[ $parse_rc -ne 0 ]]; then
        exit "$parse_rc"
    fi

    # Validate target path
    if [[ ! -d "$TARGET_PATH" ]]; then
        echo "Error: '${TARGET_PATH}' is not a valid directory" >&2
        exit "$ERR_PATH"
    fi

    # Load ignore configuration
    load_ignore_config

    # Output to terminal
    output_terminal

    # Save markdown
    save_markdown "$OUTPUT_FILE"

    exit "$ERR_OK"
}

# Entry point
main "$@"