#!/bin/bash

# ============================================================================
#  CODEBASE COLLECTOR v4 — macOS M1/ARM Safe
#  Drop into any project root, then run:   bash collect_codebase.sh
# ============================================================================

if [ -z "$BASH_VERSION" ]; then
    echo "ERROR: Run with:  bash $0"
    exit 1
fi

set -uo pipefail

# ── Config ──────────────────────────────────────────────────────────────────

OUTFILE="CODEBASE.md"
PROJECT="$(basename "$(pwd)")"
NOW="$(date "+%Y-%m-%d %H:%M:%S")"
MAX_KB=500

# Directories to skip (all levels)
SKIP_DIRS=(
    .git .svn .hg
    node_modules vendor
    venv .venv env .env
    __pycache__ .pytest_cache .mypy_cache .tox
    dist build .build DerivedData
    .next .nuxt .output .cache .parcel-cache
    coverage .nyc_output .terraform
    target Pods .cocoapods Carthage
    .gradle .idea .vscode .vs
    obj .sass-cache tmp .tmp logs
    xcuserdata .swiftpm .expo
    bower_components .serverless .amplify
    cdk.out .turbo .angular
)

# File name patterns to skip
SKIP_NAMES=(
    ".DS_Store"
    "Thumbs.db"
    "package-lock.json"
    "yarn.lock"
    "pnpm-lock.yaml"
    "Podfile.lock"
    "Gemfile.lock"
    "composer.lock"
    "Cargo.lock"
)

# Extensions to skip (binary / generated)
SKIP_EXT=(
    min.js min.css map
    pyc pyo class o a so dylib dll exe
    bin dat db sqlite sqlite3
    ico icns png jpg jpeg gif bmp svg webp avif
    mp3 mp4 wav avi mov mkv flac ogg
    pdf zip tar gz rar 7z bz2 xz
    woff woff2 ttf eot otf
    pb tfstate
    lock
)

# Extensions to include as code
CODE_EXT=(
    sh bash zsh fish
    html htm css scss sass less styl
    js jsx ts tsx mjs cjs
    vue svelte astro
    py pyi pyw pyx
    rb erb rake gemspec
    php phtml
    java kt kts scala groovy gradle
    c h cpp hpp cc cxx hxx hh
    m mm
    cs csx fs fsx vb
    swift rs go dart lua
    pl pm t
    r R Rmd
    hs lhs
    ex exs erl hrl
    clj cljs cljc edn
    zig nim v
    json jsonc json5
    yaml yml toml
    xml xsl xsd plist
    ini cfg conf properties
    md mdx rst txt adoc
    sql graphql gql proto
    tf tfvars hcl nix
    cmake mk csv
    gitignore dockerignore editorconfig
    env
)

# Exact filenames to always include
EXACT_NAMES=(
    Makefile makefile GNUmakefile CMakeLists.txt
    Dockerfile Containerfile Vagrantfile
    Rakefile Gemfile Procfile Brewfile
    Podfile Fastfile Dangerfile Jenkinsfile
    Justfile Taskfile Earthfile Caddyfile Tiltfile
    .gitignore .gitattributes .editorconfig
    .eslintrc .prettierrc .babelrc .nvmrc
    .ruby-version .python-version .tool-versions
    .env.example .env.sample .flake8 .pylintrc
    .swiftlint.yml .rubocop.yml
    setup.cfg setup.py pyproject.toml
    requirements.txt constraints.txt
    docker-compose.yml docker-compose.yaml
    tsconfig.json jsconfig.json
    webpack.config.js vite.config.ts vite.config.js
    rollup.config.js esbuild.config.js
    tailwind.config.js tailwind.config.ts
    postcss.config.js postcss.config.cjs
    next.config.js next.config.mjs next.config.ts
    nuxt.config.ts svelte.config.js astro.config.mjs
    angular.json nx.json turbo.json lerna.json
    jest.config.js jest.config.ts vitest.config.ts
    playwright.config.ts cypress.config.ts
    Cargo.toml go.mod go.sum
    Package.swift Package.resolved pubspec.yaml
    build.gradle settings.gradle pom.xml
    mix.exs rebar.config stack.yaml
    flake.nix shell.nix default.nix
    deno.json deno.jsonc bun.lockb
    nginx.conf project.pbxproj
)

# ── Temp files ──────────────────────────────────────────────────────────────

TMP_FILES="$(mktemp /tmp/cbc_files.XXXXXX)"
TMP_OUT="$(mktemp /tmp/cbc_out.XXXXXX)"
cleanup() { rm -f "$TMP_FILES" "$TMP_OUT"; }
trap cleanup EXIT

# ── Helpers ─────────────────────────────────────────────────────────────────

get_ext() {
    local name="$1"
    case "$name" in
        *.*) echo "${name##*.}" ;;
        *)   echo "" ;;
    esac
}

in_array() {
    local needle="$1"
    shift
    local item
    for item in "$@"; do
        if [ "$needle" = "$item" ]; then
            return 0
        fi
    done
    return 1
}

is_binary_file() {
    local mime
    mime="$(file --mime-encoding -- "$1" 2>/dev/null || echo "binary")"
    case "$mime" in
        *binary*) return 0 ;;
        *)        return 1 ;;
    esac
}

file_bytes() {
    wc -c < "$1" 2>/dev/null | tr -d '[:space:]'
}

file_lines() {
    wc -l < "$1" 2>/dev/null | tr -d '[:space:]'
}

human_size() {
    echo "$1" | awk '{
        if ($1 > 1048576)     printf "%.1fMB", $1/1048576
        else if ($1 > 1024)   printf "%.1fKB", $1/1024
        else                  printf "%dB", $1
    }'
}

get_lang() {
    local base ext
    base="$(basename "$1")"
    ext="$(get_ext "$base")"

    # Special filenames first
    case "$base" in
        Makefile|makefile|GNUmakefile)     echo "makefile"    ; return ;;
        Dockerfile*|Containerfile)         echo "dockerfile"  ; return ;;
        Vagrantfile|Gemfile|Rakefile)      echo "ruby"        ; return ;;
        Podfile|Fastfile|Dangerfile)       echo "ruby"        ; return ;;
        Jenkinsfile)                       echo "groovy"      ; return ;;
        Justfile)                          echo "just"        ; return ;;
        Tiltfile|Earthfile)                echo "python"      ; return ;;
        .gitignore|.dockerignore)          echo "gitignore"   ; return ;;
        .editorconfig)                     echo "ini"         ; return ;;
        .env*)                             echo "bash"        ; return ;;
        nginx.conf)                        echo "nginx"       ; return ;;
        *.config.js|*.config.ts)           echo "javascript"  ; return ;;
    esac

    # By extension
    case "$ext" in
        sh|bash|zsh|fish)            echo "bash" ;;
        py|pyi|pyw|pyx)              echo "python" ;;
        js|mjs|cjs)                  echo "javascript" ;;
        jsx)                         echo "jsx" ;;
        ts)                          echo "typescript" ;;
        tsx)                         echo "tsx" ;;
        rb|rake|gemspec|erb)         echo "ruby" ;;
        php|phtml)                   echo "php" ;;
        java)                        echo "java" ;;
        kt|kts)                      echo "kotlin" ;;
        scala)                       echo "scala" ;;
        groovy|gradle)               echo "groovy" ;;
        c|h)                         echo "c" ;;
        cpp|hpp|cc|cxx|hxx|hh)       echo "cpp" ;;
        m|mm)                        echo "objectivec" ;;
        cs|csx)                      echo "csharp" ;;
        fs|fsx)                      echo "fsharp" ;;
        vb)                          echo "vb" ;;
        swift)                       echo "swift" ;;
        rs)                          echo "rust" ;;
        go)                          echo "go" ;;
        dart)                        echo "dart" ;;
        lua)                         echo "lua" ;;
        pl|pm|t)                     echo "perl" ;;
        r|R)                         echo "r" ;;
        Rmd)                         echo "rmarkdown" ;;
        hs|lhs)                      echo "haskell" ;;
        ex|exs)                      echo "elixir" ;;
        erl|hrl)                     echo "erlang" ;;
        clj|cljs|cljc|edn)          echo "clojure" ;;
        zig)                         echo "zig" ;;
        nim)                         echo "nim" ;;
        v)                           echo "v" ;;
        html|htm)                    echo "html" ;;
        css)                         echo "css" ;;
        scss)                        echo "scss" ;;
        sass)                        echo "sass" ;;
        less)                        echo "less" ;;
        styl)                        echo "stylus" ;;
        vue)                         echo "vue" ;;
        svelte)                      echo "svelte" ;;
        astro)                       echo "astro" ;;
        json|jsonc|json5)            echo "json" ;;
        yaml|yml)                    echo "yaml" ;;
        toml)                        echo "toml" ;;
        xml|xsl|xsd|plist)          echo "xml" ;;
        ini|cfg|conf|properties)     echo "ini" ;;
        env)                         echo "bash" ;;
        md|mdx)                      echo "markdown" ;;
        rst)                         echo "rst" ;;
        txt|adoc)                    echo "text" ;;
        sql)                         echo "sql" ;;
        graphql|gql)                 echo "graphql" ;;
        proto)                       echo "protobuf" ;;
        tf|tfvars|hcl)               echo "hcl" ;;
        nix)                         echo "nix" ;;
        cmake)                       echo "cmake" ;;
        *)                           echo "" ;;
    esac
}

# ── Check: should we process this file? ─────────────────────────────────────

check_file() {
    local filepath="$1"
    local filename ext bytes kb

    filename="$(basename "$filepath")"
    ext="$(get_ext "$filename")"

    # Skip self and output
    [ "$filename" = "$OUTFILE" ] && return 1
    [ "$filename" = "collect_codebase.sh" ] && return 1

    # Skip by exact name
    if in_array "$filename" "${SKIP_NAMES[@]}"; then
        return 1
    fi

    # Skip by extension
    if [ -n "$ext" ] && in_array "$ext" "${SKIP_EXT[@]}"; then
        return 1
    fi

    # Check size
    bytes="$(file_bytes "$filepath")"
    bytes="${bytes:-0}"
    kb=$((bytes / 1024))
    if [ "$kb" -gt "$MAX_KB" ]; then
        return 1
    fi

    # Must match code extension OR exact filename
    local match=false
    if in_array "$filename" "${EXACT_NAMES[@]}"; then
        match=true
    fi
    if [ "$match" = false ] && [ -n "$ext" ] && in_array "$ext" "${CODE_EXT[@]}"; then
        match=true
    fi
    if [ "$match" = false ]; then
        return 1
    fi

    # Skip binary
    if is_binary_file "$filepath"; then
        return 2
    fi

    return 0
}

# ── Build find command to get all files ─────────────────────────────────────

find_all_files() {
    # Build prune expression for directories
    local prune_args=""
    local d
    for d in "${SKIP_DIRS[@]}"; do
        if [ -n "$prune_args" ]; then
            prune_args="$prune_args -o -name $d"
        else
            prune_args="-name $d"
        fi
    done

    # Find all regular files, pruning skip dirs, sorted
    eval "find . \\( $prune_args \\) -prune -o -type f -print" 2>/dev/null | sort
}

# ── Build directory tree ────────────────────────────────────────────────────

build_tree() {
    if command -v tree >/dev/null 2>&1; then
        local ignore
        ignore="$(IFS='|'; echo "${SKIP_DIRS[*]}")"
        tree -a -I "$ignore" --charset utf-8 --dirsfirst -F 2>/dev/null
    else
        echo "."
        # Simple fallback using find
        local prune_args=""
        local d
        for d in "${SKIP_DIRS[@]}"; do
            if [ -n "$prune_args" ]; then
                prune_args="$prune_args -o -name $d"
            else
                prune_args="-name $d"
            fi
        done
        eval "find . \\( $prune_args \\) -prune -o -print" 2>/dev/null \
            | sed 1d \
            | sort \
            | awk '{
                # Count depth by slashes
                n = gsub(/\//, "/")
                indent = ""
                for (i = 0; i < n; i++) indent = indent "    "
                # Get basename
                split($0, parts, "/")
                name = parts[length(parts)]
                print indent "├── " name
            }' \
            | head -800
    fi
}

# ── Banner ──────────────────────────────────────────────────────────────────

echo ""
echo "╔══════════════════════════════════════════════════════╗"
echo "║           CODEBASE COLLECTOR v4                      ║"
echo "╚══════════════════════════════════════════════════════════╝"
echo ""
echo "  Project:  $PROJECT"
echo "  Output:   $OUTFILE"
echo ""
echo "  Scanning all folders and subfolders..."
echo ""

# ── Collect all file paths ──────────────────────────────────────────────────

find_all_files > "$TMP_FILES"

total_found="$(wc -l < "$TMP_FILES" | tr -d '[:space:]')"
echo "  Found $total_found files total. Filtering..."

# ── Write markdown ──────────────────────────────────────────────────────────

file_count=0
skipped_count=0
binary_count=0
total_lines=0

# --- Header ---
cat > "$TMP_OUT" <<HEADER
# Codebase: \`${PROJECT}\`

> Generated on ${NOW}

---

## Directory Structure

\`\`\`
HEADER

build_tree >> "$TMP_OUT"

echo '```' >> "$TMP_OUT"
echo "" >> "$TMP_OUT"
echo "---" >> "$TMP_OUT"
echo "" >> "$TMP_OUT"
echo "## File Contents" >> "$TMP_OUT"
echo "" >> "$TMP_OUT"

# --- Each file ---
while IFS= read -r fpath; do

    # Strip leading ./
    fpath="${fpath#./}"

    # Skip empty lines
    [ -z "$fpath" ] && continue

    # Check file
    check_file "$fpath"
    result=$?

    if [ $result -eq 1 ]; then
        skipped_count=$((skipped_count + 1))
        continue
    fi
    if [ $result -eq 2 ]; then
        binary_count=$((binary_count + 1))
        continue
    fi

    # Get metadata
    lang="$(get_lang "$fpath")"
    lc="$(file_lines "$fpath")"
    lc="${lc:-0}"
    total_lines=$((total_lines + lc))
    bytes="$(file_bytes "$fpath")"
    bytes="${bytes:-0}"
    sz="$(human_size "$bytes")"

    # Write separator and header
    echo "" >> "$TMP_OUT"
    echo "---" >> "$TMP_OUT"
    echo "" >> "$TMP_OUT"
    echo "### \`$fpath\`" >> "$TMP_OUT"
    echo "" >> "$TMP_OUT"
    echo "> Lines: $lc | Size: $sz" >> "$TMP_OUT"
    echo "" >> "$TMP_OUT"

    # Opening code fence
    echo "\`\`\`$lang" >> "$TMP_OUT"

    # Dump file contents safely with cat
    cat -- "$fpath" >> "$TMP_OUT" 2>/dev/null

    # Ensure trailing newline before closing fence
    echo "" >> "$TMP_OUT"

    # Closing code fence
    echo '```' >> "$TMP_OUT"

    file_count=$((file_count + 1))

    # Progress indicator
    if [ $((file_count % 10)) -eq 0 ]; then
        echo "  ... processed $file_count files" >&2
    fi

done < "$TMP_FILES"

# --- Footer ---
echo "" >> "$TMP_OUT"
echo "---" >> "$TMP_OUT"
echo "" >> "$TMP_OUT"
echo "## Summary" >> "$TMP_OUT"
echo "" >> "$TMP_OUT"
echo "| Metric | Value |" >> "$TMP_OUT"
echo "|--------|-------|" >> "$TMP_OUT"
echo "| Files included | $file_count |" >> "$TMP_OUT"
echo "| Files skipped | $skipped_count |" >> "$TMP_OUT"
echo "| Binary skipped | $binary_count |" >> "$TMP_OUT"
echo "| Total lines | $total_lines |" >> "$TMP_OUT"
echo "| Generated | $NOW |" >> "$TMP_OUT"
echo "" >> "$TMP_OUT"
echo "---" >> "$TMP_OUT"

# ── Move output ─────────────────────────────────────────────────────────────

mv -f "$TMP_OUT" "$OUTFILE"

# ── Results ─────────────────────────────────────────────────────────────────

out_bytes="$(file_bytes "$OUTFILE")"
out_bytes="${out_bytes:-0}"
out_sz="$(human_size "$out_bytes")"

echo ""
echo "  ================================================"
echo "    DONE"
echo "  ================================================"
echo ""
echo "    Output:          $OUTFILE"
echo "    Size:            $out_sz"
echo "    Files collected: $file_count"
echo "    Files skipped:   $skipped_count"
echo "    Binary skipped:  $binary_count"
echo "    Total lines:     $total_lines"
echo ""
echo "  ================================================"
echo ""

if [ "$out_bytes" -gt 10485760 ]; then
    echo "  WARNING: Output exceeds 10MB."
    echo "  Add more dirs to SKIP_DIRS or lower MAX_KB."
    echo ""
fi
