codegen-document

#!/usr/bin/env bash
# codegen-document — generate usage_rules docs for Elixir hex dependencies
#
# Usage:
#   codegen-document              # whole-lock mode: generate docs for all deps in mix.lock
#   codegen-document <dep_name>   # single-dep mode: generate doc for one dep
#
# Output: $CODEGEN_DIR/shared/usage_rules/<dep>-<version>.md
#         (frameworks split into multiple topic files)
#
# Requires: mix.lock in cwd (or git-repo root). codegen-call on PATH.
# Exit codes: 0 = ok, 1 = error (no mix.lock / dep not found)

set -euo pipefail

# Derive CODEGEN_DIR from this script's own location (works symlinked via PATH)
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CODEGEN_DIR="$SCRIPT_DIR"

USAGE_RULES_DIR="$CODEGEN_DIR/shared/usage_rules"
SYSTEM_PROMPT_FILE="$CODEGEN_DIR/harnesses/claude/usage-rules-system-prompt.md"
REGEN_INDEX_SCRIPT="$CODEGEN_DIR/templates/generator/regenerate_usage_rules_index.py"
MODEL="haiku"

# Track background jobs for cleanup
declare -a BG_PIDS=()

cleanup() {
    echo ""
    echo "Interrupted! Killing background jobs..."
    if [[ ${#BG_PIDS[@]} -gt 0 ]]; then
        for pid in "${BG_PIDS[@]}"; do
            kill "$pid" 2>/dev/null || true
        done
    fi
    exit 130
}

trap cleanup INT TERM

# Ensure usage rules directory exists
mkdir -p "$USAGE_RULES_DIR"

# Built-in Elixir/Erlang applications that don't need documentation
BUILTIN_DEPS="logger crypto ssl public_key inets mnesia runtime_tools observer sasl tools eex iex mix ex_unit"

# Function to check if dependency is built-in
is_builtin() {
    local dep="$1"
    echo "$BUILTIN_DEPS" | grep -qw "$dep"
}

# Find the mix.lock file from cwd or git root
find_mix_lock() {
    if [ -f "./mix.lock" ]; then
        echo "./mix.lock"
        return 0
    fi
    # Fallback: git repo root
    local git_root
    git_root="$(git rev-parse --show-toplevel 2>/dev/null)" || true
    if [ -n "$git_root" ] && [ -f "$git_root/mix.lock" ]; then
        echo "$git_root/mix.lock"
        return 0
    fi
    return 1
}

# Function to extract dependencies with versions from mix.lock
# Usage: extract_dependencies <mix_lock_file>
extract_dependencies() {
    local lock_file="$1"

    if [ ! -f "$lock_file" ]; then
        echo "Error: mix.lock not found at $lock_file" >&2
        exit 1
    fi

    # Hex format: "dep_name": {:hex, :dep_name, "1.2.3", ...}
    # Git format: "dep_name": {:git, "url", "sha", [tag: "v1.2.3", ...]}
    while IFS= read -r line; do
        # Extract dep name from quoted key at line start
        local dep_name
        dep_name=$(echo "$line" | grep -oE '^  "[a-z_0-9]+"' | tr -d ' "') || true
        [ -z "$dep_name" ] && continue

        is_builtin "$dep_name" && continue

        # Reset all branch-locals at top of each iteration
        local version="" repo_url="" tree_ref=""

        # Check if it's a hex package
        if echo "$line" | grep -q ":hex,"; then
            version=$(echo "$line" | grep -oE '"[0-9]+\.[0-9]+(\.[0-9]+(-[a-z0-9.]+)?)?"' | head -1 | tr -d '"')
            [ -z "$version" ] && version="~no_version~"
            echo "${dep_name}:${version}"
        elif echo "$line" | grep -q ":git,"; then
            local tag
            tag=$(echo "$line" | grep -oE 'tag: "[^"]*"' | cut -d'"' -f2) || true
            if [ -n "$tag" ]; then
                version=$(echo "$tag" | sed 's/^v//')
            else
                local sha
                sha=$(echo "$line" | grep -oE '"[a-f0-9]{40}"' | head -1 | tr -d '"' | cut -c1-7) || true
                [ -n "$sha" ] && version="git-$sha"
            fi
            [ -z "$version" ] && version="~no_version~"
            # Capture GitHub repo path (owner/repo) and full commit SHA
            repo_url=$(echo "$line" | grep -oE 'github\.com/[^"]+\.git' | sed 's#github\.com/##; s#\.git##') || true
            tree_ref=$(echo "$line" | grep -oE '"[a-f0-9]{40}"' | head -1 | tr -d '"') || true
            echo "${dep_name}:${version}|${repo_url}|${tree_ref}"
        else
            version="~no_version~"
            echo "${dep_name}:${version}"
        fi
    done <"$lock_file" | sort -u
}

# Function to check if usage rule needs updating (file absent = needs generation)
usage_rule_needs_update() {
    local dep_name="$1"
    local version="$2"
    local version_safe
    version_safe=$(echo "$version" | sed 's/[^a-zA-Z0-9.]/-/g')
    local file="$USAGE_RULES_DIR/${dep_name}-${version_safe}.md"
    [ ! -f "$file" ]
}

# Function to generate usage rule from hexdocs via LLM
# $1 = dep_name, $2 = version, $3 = full record (dep_name:version|repo_url|tree_ref) — optional
generate_usage_rule() {
    local dep_name="$1"
    local version="$2"
    local record="${3:-}"
    local version_safe
    version_safe=$(echo "$version" | sed 's/[^a-zA-Z0-9.]/-/g')
    local output_file="$USAGE_RULES_DIR/${dep_name}-${version_safe}.md"

    local version_display="$version"
    [ "$version" = "~no_version~" ] && version_display="latest"

    # Parse git coords from record (pipe-delimited tail; absent for hex deps)
    local repo_url="" tree_ref=""
    if [ -n "$record" ]; then
        repo_url=$(echo "$record" | cut -d'|' -f2)
        tree_ref=$(echo "$record" | cut -d'|' -f3)
    fi

    local hexdocs_url
    if [ -n "$repo_url" ] && [ -n "$tree_ref" ]; then
        # git dep with captured repo — fetch from fork at locked commit
        hexdocs_url="https://github.com/${repo_url}/tree/${tree_ref}"
    elif echo "$version" | grep -qE '^[0-9]+\.[0-9]+\.[0-9]'; then
        # hex dep with real semver — versioned hexdocs
        hexdocs_url="https://hexdocs.pm/${dep_name}/${version}/"
    else
        # fallback — unversioned hexdocs
        hexdocs_url="https://hexdocs.pm/${dep_name}/"
    fi

    # Check if this is a framework that needs multi-file splitting
    local is_framework=false
    case "$dep_name" in
    phoenix | ecto | phoenix_live_view | absinthe | ash | nerves | broadway | oban | membrane)
        is_framework=true
        ;;
    esac

    local prompt_file
    prompt_file=$(mktemp)
    local temp_output
    temp_output=$(mktemp)
    local today
    today=$(date '+%Y-%m-%d')

    # Build per-dep prompt based on framework/single branch
    if [ "$is_framework" = true ]; then
        cat >"$prompt_file" <<PROMPT_EOF
Visit ${hexdocs_url} and analyze the documentation to create usage rules.

CRITICAL: You have Write tool access. Use it to create files directly in ${USAGE_RULES_DIR}.

This is a FRAMEWORK library. Split documentation into MULTIPLE SEPARATE FILES:

1. MAIN INDEX FILE: ${USAGE_RULES_DIR}/${dep_name}-${version_safe}.md
   Content:
   # ${dep_name}

   [2-3 paragraph overview of what the library does]

   ## Quick Start
   [Basic installation and minimal example]

   ## Documentation Sections
   - [Topic 1 Name](${dep_name}-${version_safe}-topic1.md)
   - [Topic 2 Name](${dep_name}-${version_safe}-topic2.md)
   - [Topic 3 Name](${dep_name}-${version_safe}-topic3.md)
   ... (5-8 major topics based on guide structure)

   ---
   **Version:** ${version_display}
   **Source:** [hexdocs.pm/${dep_name}](${hexdocs_url})
   **Generated:** ${today}

2. SEPARATE TOPIC FILES: ${USAGE_RULES_DIR}/${dep_name}-${version_safe}-[topic].md
   One file per major topic (5-8 files total).

   Each file structure:
   # ${dep_name} - [Topic Name]

   [Comprehensive content extracted from hexdocs guides]
   [Include code examples, patterns, best practices]
   [100-250 lines per file]

   ---
   [← Back to main](${dep_name}-${version_safe}.md)
   **Version:** ${version_display}

MANDATORY REQUIREMENTS:
- Use Write tool to create files directly in ${USAGE_RULES_DIR}
- NO conversational text before/after - ONLY create the files
- Extract actual content from hexdocs guides (not API docs)
- Determine 5-8 logical topic splits based on guide structure
- Each section file should be comprehensive (100-250 lines)
- After creating files, output ONLY: "Created: [list of filenames]"
PROMPT_EOF
    else
        cat >"$prompt_file" <<PROMPT_EOF
Visit ${hexdocs_url} and create usage documentation.

CRITICAL: You have Write tool access. Use it to create files directly in ${USAGE_RULES_DIR}.

Create ONE COMPREHENSIVE FILE: ${USAGE_RULES_DIR}/${dep_name}-${version_safe}.md

Structure:
# ${dep_name}
[Overview]
## Quick Start
[Installation and basic usage]
## Core Concepts
[Main features and patterns]
## Configuration
[How to configure]
## Best Practices
[Common patterns and recommendations]
---
**Version:** ${version_display}
**Source:** [hexdocs.pm/${dep_name}](${hexdocs_url})
**Generated:** ${today}

MANDATORY REQUIREMENTS:
- Use Write tool to create the file directly in ${USAGE_RULES_DIR}
- NO conversational text before/after - ONLY create the file
- Extract actual content from hexdocs guides (not API docs)
- Keep comprehensive but concise (100-300 lines)
- After creating file, output ONLY: "Created: [filename]"
PROMPT_EOF
    fi

    local per_dep_prompt
    per_dep_prompt=$(cat "$prompt_file")
    rm -f "$prompt_file"

    echo "Generating usage rule for ${dep_name} (${version_display})..."

    # Try up to 3 times if LLM generation fails
    local max_retries=3
    local retry=0
    local success=false

    while [ $retry -lt $max_retries ] && [ "$success" = false ]; do
        if codegen-call --harness claude_code --model "$MODEL" --effort low \
            --system-prompt @"$SYSTEM_PROMPT_FILE" \
            --allowed-tools "Write,WebFetch" \
            -- "$per_dep_prompt" >"$temp_output" 2>/dev/null; then
            # Check if main file was created
            if [ -f "$output_file" ]; then
                # Count additional section files (for large libraries)
                local section_count=0
                while IFS= read -r _; do
                    section_count=$((section_count + 1))
                done < <(ls "${USAGE_RULES_DIR}/${dep_name}-${version_safe}"-*.md 2>/dev/null | grep -v "^${output_file}$" || true)

                if [ "$section_count" -gt 0 ]; then
                    echo "Generated: ${output_file} + ${section_count} section files (version: ${version_display})"
                else
                    echo "Generated: ${output_file} (version: ${version_display})"
                fi
                success=true
            else
                retry=$((retry + 1))
                if [ $retry -lt $max_retries ]; then
                    echo "File not created for ${dep_name}, retrying (${retry}/${max_retries})..." >&2
                    sleep 2
                else
                    echo "Failed to generate usage rule for ${dep_name} after ${max_retries} attempts" >&2
                    rm -f "$temp_output"
                    return 1
                fi
            fi
        else
            retry=$((retry + 1))
            if [ $retry -lt $max_retries ]; then
                echo "LLM call failed for ${dep_name}, retrying (${retry}/${max_retries})..." >&2
                sleep 2
            else
                echo "Failed to generate usage rule for ${dep_name} after ${max_retries} attempts" >&2
                rm -f "$temp_output"
                return 1
            fi
        fi
    done

    rm -f "$temp_output"
}

# Function to regenerate shared/usage_rules/INDEX.md from what is on disk.
# Non-fatal: a regeneration failure prints a warning but never aborts
# codegen-document (the per-dep .md files are already written by this point).
regenerate_index() {
    if [ ! -f "$REGEN_INDEX_SCRIPT" ]; then
        echo "Warning: $REGEN_INDEX_SCRIPT not found — skipping INDEX.md regeneration" >&2
        return 0
    fi
    if python3 "$REGEN_INDEX_SCRIPT" --usage-rules-dir "$USAGE_RULES_DIR" --output "$USAGE_RULES_DIR/INDEX.md"; then
        return 0
    else
        echo "Warning: INDEX.md regeneration failed — corpus docs were written but INDEX.md may be stale" >&2
        return 0
    fi
}

# Main function — whole-lock mode
main() {
    echo "codegen-document — Usage Rules Generator"
    echo "========================================="
    echo "Model: ${MODEL}"
    echo ""

    local lock_file
    if ! lock_file=$(find_mix_lock); then
        echo "Error: mix.lock not found in cwd or git root. Run from a Phoenix/Elixir project." >&2
        exit 1
    fi

    echo "Lock file: ${lock_file}"
    echo "Output dir: ${USAGE_RULES_DIR}"
    echo ""

    local dependencies
    dependencies=$(extract_dependencies "$lock_file")

    if [ -z "$dependencies" ]; then
        echo "No dependencies found in mix.lock"
        exit 0
    fi

    echo "Dependencies found:"
    echo "$dependencies" | sed 's/^/  - /'
    echo ""

    # Process with throttled parallelism (max 5 concurrent)
    local new_rules=0
    local existing_rules=0
    local failed_rules=0
    local deps_to_generate=()
    local pids=()
    local temp_dir
    temp_dir=$(mktemp -d)

    # First pass: identify which deps need generation
    while IFS= read -r dep_with_version; do
        [ -z "$dep_with_version" ] && continue
        local dep_name dep_version
        dep_name=$(echo "$dep_with_version" | cut -d: -f1)
        dep_version=$(echo "$dep_with_version" | cut -d: -f2 | cut -d'|' -f1)

        if usage_rule_needs_update "$dep_name" "$dep_version"; then
            deps_to_generate+=("$dep_with_version")
        else
            local vdisplay="$dep_version"
            [ "$dep_version" = "~no_version~" ] && vdisplay="latest"
            echo "Skipping ${dep_name} (${vdisplay}) — already exists"
            existing_rules=$((existing_rules + 1))
        fi
    done <<<"$dependencies"

    # Second pass: generate with throttling (max 5 concurrent)
    local max_concurrent=5
    local total_to_generate="${#deps_to_generate[@]}"
    local current_index=0

    while [ "$current_index" -lt "$total_to_generate" ] || [ "${#pids[@]}" -gt 0 ]; do
        # Start new jobs up to max_concurrent
        local current_pids="${#pids[@]}"
        while [ "$current_pids" -lt "$max_concurrent" ] && [ "$current_index" -lt "$total_to_generate" ]; do
            local dep_with_version="${deps_to_generate[$current_index]}"
            local dep_name dep_version
            dep_name=$(echo "$dep_with_version" | cut -d: -f1)
            dep_version=$(echo "$dep_with_version" | cut -d: -f2 | cut -d'|' -f1)
            (
                if generate_usage_rule "$dep_name" "$dep_version" "$dep_with_version"; then
                    echo "success" >"$temp_dir/${dep_name}.status"
                else
                    echo "failed" >"$temp_dir/${dep_name}.status"
                fi
            ) &
            local new_pid=$!
            pids+=("$new_pid")
            BG_PIDS+=("$new_pid")
            current_index=$((current_index + 1))
            current_pids=$((current_pids + 1))
        done

        # Wait for at least one job to complete
        if [ "${#pids[@]}" -gt 0 ]; then
            wait -n "${pids[@]}" 2>/dev/null || true
            # Remove completed PIDs
            local new_pids=()
            local _pid
            for _pid in "${pids[@]}"; do
                if kill -0 "$_pid" 2>/dev/null; then
                    new_pids+=("$_pid")
                fi
            done
            if [[ ${#new_pids[@]} -gt 0 ]]; then
                pids=("${new_pids[@]}")
            else
                pids=()
            fi
        fi
    done

    # Count results
    if [[ ${#deps_to_generate[@]} -gt 0 ]]; then
        for dep_with_version in "${deps_to_generate[@]}"; do
            local dep_name
            dep_name=$(echo "$dep_with_version" | cut -d: -f1)
            if [ -f "$temp_dir/${dep_name}.status" ]; then
                local status
                status=$(cat "$temp_dir/${dep_name}.status")
                if [ "$status" = "success" ]; then
                    new_rules=$((new_rules + 1))
                else
                    failed_rules=$((failed_rules + 1))
                fi
            else
                failed_rules=$((failed_rules + 1))
            fi
        done
    fi

    rm -rf "$temp_dir"

    echo ""
    echo "Summary:"
    echo "  New usage rules generated: ${new_rules}"
    echo "  Existing usage rules found: ${existing_rules}"
    echo "  Failed generations: ${failed_rules}"
    echo "  Total processed: $((new_rules + existing_rules + failed_rules))"

    regenerate_index

    if [ "$failed_rules" -gt 0 ]; then
        echo ""
        echo "Some dependencies failed after 3 retries. Run again to retry, or check hexdocs.pm." >&2
        return 1
    fi

    return 0
}

# ─── Argument parsing ─────────────────────────────────────────────────────────
MODEL_OVERRIDE=""
POSITIONAL_DEP=""

while [[ $# -gt 0 ]]; do
    case $1 in
    --model | -m)
        MODEL_OVERRIDE="$2"
        shift 2
        ;;
    --help | -h)
        echo "Usage: codegen-document [options] [dep_name]"
        echo ""
        echo "  dep_name            Single dep mode: generate doc for one dep (looks up version from mix.lock)"
        echo "  (no dep_name)       Whole-lock mode: generate docs for all deps in mix.lock"
        echo ""
        echo "Options:"
        echo "  --model, -m <model>   AI model to use (haiku/sonnet/opus, default: haiku)"
        echo "  --help, -h            Show this help"
        echo ""
        echo "Examples:"
        echo "  codegen-document              # generate all missing docs"
        echo "  codegen-document jason        # generate/regenerate doc for jason"
        echo "  codegen-document ecto --model sonnet"
        echo ""
        echo "After generation, shared/usage_rules/INDEX.md is regenerated automatically"
        echo "(re-citing the max version present for each already-tracked dep)."
        exit 0
        ;;
    -*)
        echo "Unknown option: $1" >&2
        echo "Use --help for usage information" >&2
        exit 1
        ;;
    *)
        if [ -n "$POSITIONAL_DEP" ]; then
            echo "Error: unexpected argument '$1' (dep_name already set to '$POSITIONAL_DEP')" >&2
            exit 1
        fi
        POSITIONAL_DEP="$1"
        shift
        ;;
    esac
done

# Apply model override
if [ -n "$MODEL_OVERRIDE" ]; then
    case "$MODEL_OVERRIDE" in
    haiku | sonnet | opus)
        MODEL="$MODEL_OVERRIDE"
        ;;
    *)
        echo "Invalid model '$MODEL_OVERRIDE'. Valid models: haiku, sonnet, opus." >&2
        exit 1
        ;;
    esac
fi

# ─── Single-dep mode ──────────────────────────────────────────────────────────
if [ -n "$POSITIONAL_DEP" ]; then
    dep_arg="$POSITIONAL_DEP"

    lock_file=""
    lock_file=$(find_mix_lock) || true
    if [ -z "$lock_file" ]; then
        echo "Error: mix.lock not found in cwd or git root." >&2
        exit 1
    fi

    # Look up version from mix.lock
    dep_line=$(grep "^  \"${dep_arg}\":" "$lock_file") || true
    if [ -z "$dep_line" ]; then
        echo "Error: '${dep_arg}' not found in ${lock_file}." >&2
        exit 1
    fi

    dep_version=""
    dep_repo_url=""
    dep_tree_ref=""
    if echo "$dep_line" | grep -q ":hex,"; then
        dep_version=$(echo "$dep_line" | grep -oE '"[0-9]+\.[0-9]+(\.[0-9]+(-[a-z0-9.]+)?)?"' | head -1 | tr -d '"')
    elif echo "$dep_line" | grep -q ":git,"; then
        local_tag=$(echo "$dep_line" | grep -oE 'tag: "[^"]*"' | cut -d'"' -f2) || true
        if [ -n "${local_tag:-}" ]; then
            dep_version=$(echo "$local_tag" | sed 's/^v//')
        else
            local_sha=$(echo "$dep_line" | grep -oE '"[a-f0-9]{40}"' | head -1 | tr -d '"' | cut -c1-7) || true
            [ -n "${local_sha:-}" ] && dep_version="git-${local_sha}"
        fi
        dep_repo_url=$(echo "$dep_line" | grep -oE 'github\.com/[^"]+\.git' | sed 's#github\.com/##; s#\.git##') || true
        dep_tree_ref=$(echo "$dep_line" | grep -oE '"[a-f0-9]{40}"' | head -1 | tr -d '"') || true
    fi
    [ -z "$dep_version" ] && dep_version="~no_version~"

    dep_record="${dep_arg}:${dep_version}|${dep_repo_url}|${dep_tree_ref}"
    generate_usage_rule "$dep_arg" "$dep_version" "$dep_record"
    gen_rc=$?
    regenerate_index
    exit $gen_rc
fi

# ─── Whole-lock mode ──────────────────────────────────────────────────────────
main