#!/usr/bin/env bash
#
# ONEX Node Regenerator - Claude Code Skill
#
# Regenerates an existing ONEX node by:
# 1. Extracting description from README.md (if exists)
# 2. Using Z.ai to analyze code and extract prompt (if no README)
# 3. Calling the generate skill with the extracted prompt
#
# Usage:
#   ${CLAUDE_PLUGIN_ROOT}/skills/generate-node/regenerate <NODE_DIR> [OPTIONS]
#
# Examples:
#   ${CLAUDE_PLUGIN_ROOT}/skills/generate-node/regenerate src/omninode_bridge/nodes/llm_effect/v1_0_0/llm_effect_llm
#   ${CLAUDE_PLUGIN_ROOT}/skills/generate-node/regenerate nodes/my_custom_node --output-dir ./regenerated
#

set -euo pipefail

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
MAGENTA='\033[0;35m'
NC='\033[0m' # No Color

# Repository path
REPO_PATH="/Volumes/PRO-G40/Code/omninode_bridge"

# Function to print colored messages
print_error() {
    echo -e "${RED}❌ Error: $1${NC}" >&2
}

print_success() {
    echo -e "${GREEN}✅ $1${NC}"
}

print_info() {
    echo -e "${BLUE}ℹ️  $1${NC}"
}

print_warning() {
    echo -e "${YELLOW}⚠️  $1${NC}"
}

print_step() {
    echo -e "${MAGENTA}🔧 $1${NC}"
}

# Function to check prerequisites
check_prerequisites() {
    # Check if repository exists
    if [[ ! -d "$REPO_PATH" ]]; then
        print_error "Repository not found: $REPO_PATH"
        exit 1
    fi

    # Check if jq is available (for JSON parsing)
    if ! command -v jq &> /dev/null; then
        print_error "jq not found. Please install jq first."
        print_info "Install with: brew install jq"
        exit 1
    fi

    # Check if .env file exists and has ZAI_API_KEY
    if [[ ! -f "$REPO_PATH/.env" ]]; then
        print_error ".env file not found in $REPO_PATH"
        exit 1
    fi

    # Source .env and check for ZAI_API_KEY
    set -a
    source "$REPO_PATH/.env"
    set +a

    if [[ -z "${ZAI_API_KEY:-}" ]]; then
        print_error "ZAI_API_KEY not set in .env file"
        print_info "Add ZAI_API_KEY=your-key-here to $REPO_PATH/.env"
        exit 1
    fi
}

# Function to extract prompt from README
extract_prompt_from_readme() {
    local readme_path="$1"

    print_step "Extracting description from README.md..."

    # Extract first paragraph after title or first non-empty paragraph
    local description=$(grep -v "^#" "$readme_path" | \
                       grep -v "^$" | \
                       grep -v "^-" | \
                       head -3 | \
                       tr '\n' ' ' | \
                       sed 's/  */ /g' | \
                       sed 's/^ *//; s/ *$//')

    if [[ -n "$description" ]]; then
        echo "$description"
        return 0
    fi

    # Fallback: extract from title
    local title=$(grep "^# " "$readme_path" | head -1 | sed 's/^# //')
    if [[ -n "$title" ]]; then
        echo "Regenerate: $title"
        return 0
    fi

    return 1
}

# Function to extract prompt using Z.ai
extract_prompt_with_zai() {
    local node_dir="$1"

    print_step "No README found. Analyzing code with Z.ai..."

    # Collect all Python files
    local python_files=$(find "$node_dir" -name "*.py" -type f 2>/dev/null || true)

    if [[ -z "$python_files" ]]; then
        print_error "No Python files found in $node_dir"
        exit 1
    fi

    # Concatenate code from all Python files
    local code_content=""
    local file_count=0
    for py_file in $python_files; do
        if [[ -f "$py_file" ]]; then
            code_content+="# File: $(basename "$py_file")\n"
            code_content+="$(cat "$py_file")\n\n"
            ((file_count++))
        fi
    done

    print_info "Collected code from $file_count Python files"

    # Prepare Z.ai API request
    local api_url="${ZAI_ENDPOINT:-https://api.z.ai/api/anthropic}"
    local model="${ZAI_MODEL:-glm-4-flash}"

    local system_prompt="You are a code analysis expert. Analyze the provided ONEX node code and extract a concise generation prompt that describes what this node does. Focus on:
- Node type (Effect, Compute, Reducer, Orchestrator)
- Primary functionality and operations
- Key integrations and dependencies
- Important features (circuit breaker, retry, metrics, etc.)

Return ONLY the prompt text (1-2 sentences), no additional explanation."

    local user_prompt="Analyze this ONEX node code and create a generation prompt:\n\n$code_content"

    # Create JSON payload
    local json_payload=$(jq -n \
        --arg model "$model" \
        --arg system "$system_prompt" \
        --arg user "$user_prompt" \
        '{
            model: $model,
            max_tokens: 500,
            temperature: 0.3,
            messages: [
                {role: "system", content: $system},
                {role: "user", content: $user}
            ]
        }')

    print_info "Calling Z.ai API (model: $model)..."

    # Call Z.ai API
    local response=$(curl -s -X POST "$api_url/v1/messages" \
        -H "Content-Type: application/json" \
        -H "x-api-key: $ZAI_API_KEY" \
        -H "anthropic-version: 2023-06-01" \
        -d "$json_payload")

    # Extract prompt from response
    local extracted_prompt=$(echo "$response" | jq -r '.content[0].text // empty' 2>/dev/null || echo "")

    if [[ -z "$extracted_prompt" ]]; then
        # Check for error in response
        local error_msg=$(echo "$response" | jq -r '.error.message // "Unknown error"' 2>/dev/null || echo "Unknown error")
        print_error "Z.ai API call failed: $error_msg"
        print_info "Response: $response"
        exit 1
    fi

    # Clean up the prompt (remove quotes, trim)
    extracted_prompt=$(echo "$extracted_prompt" | sed 's/^["'\'']*//; s/["'\'']*$//' | xargs)

    print_success "Extracted prompt from code analysis"
    echo "$extracted_prompt"
}

# Main execution
main() {
    if [[ $# -eq 0 ]]; then
        print_error "No arguments provided"
        echo ""
        echo "Usage: $0 <NODE_DIR> [OPTIONS]"
        echo ""
        echo "Arguments:"
        echo "  NODE_DIR                        Path to node directory (absolute or relative)"
        echo ""
        echo "Options:"
        echo "  --output-dir TEXT               Output directory (default: ./regenerated_nodes)"
        echo "  --node-type [effect|orchestrator|reducer|compute]"
        echo "  --interactive                   Enable interactive checkpoints"
        echo "  --enable-intelligence           Use RAG intelligence (default: enabled)"
        echo "  --disable-intelligence          Disable RAG intelligence"
        echo "  --enable-quorum                 Use AI quorum validation"
        echo "  --timeout INTEGER               Timeout in seconds (default: 300)"
        echo ""
        echo "Examples:"
        echo "  $0 src/omninode_bridge/nodes/llm_effect/v1_0_0/llm_effect_llm"
        echo "  $0 nodes/my_node --output-dir ./regenerated"
        echo "  $0 ../other_repo/nodes/custom_node --interactive"
        echo ""
        exit 1
    fi

    print_info "ONEX Node Regenerator - Starting..."
    echo ""

    # Check prerequisites
    check_prerequisites

    # Extract node directory (first argument)
    local node_dir="$1"
    shift

    # Resolve absolute path
    if [[ "$node_dir" != /* ]]; then
        # Relative path - resolve from current directory
        node_dir="$(cd "$(dirname "$node_dir")" && pwd)/$(basename "$node_dir")"
    fi

    # Validate node directory exists
    if [[ ! -d "$node_dir" ]]; then
        print_error "Node directory not found: $node_dir"
        exit 1
    fi

    print_info "Node directory: $node_dir"

    # Extract node name from path
    local node_name=$(basename "$node_dir")
    print_info "Node name: $node_name"
    echo ""

    # Try to extract prompt
    local prompt=""
    local readme_path="$node_dir/README.md"

    if [[ -f "$readme_path" ]]; then
        print_step "Found README.md in node directory"
        prompt=$(extract_prompt_from_readme "$readme_path")

        if [[ -n "$prompt" ]]; then
            print_success "Extracted prompt from README.md"
            echo ""
            print_info "Prompt: \"$prompt\""
            echo ""
        else
            print_warning "README.md exists but couldn't extract description"
            print_step "Falling back to Z.ai code analysis..."
            prompt=$(extract_prompt_with_zai "$node_dir")
        fi
    else
        print_warning "No README.md found in node directory"
        prompt=$(extract_prompt_with_zai "$node_dir")
    fi

    # Validate extracted prompt
    if [[ -z "$prompt" ]]; then
        print_error "Failed to extract generation prompt"
        exit 1
    fi

    echo ""
    print_info "Final prompt: \"$prompt\""
    echo ""

    # Call generate skill with extracted prompt
    print_step "Calling generate skill..."
    echo ""

    # Default output dir to regenerated_nodes if not specified
    local has_output_dir=false
    for arg in "$@"; do
        if [[ "$arg" == "--output-dir" ]]; then
            has_output_dir=true
            break
        fi
    done

    if [[ "$has_output_dir" == false ]]; then
        set -- "$@" --output-dir "./regenerated_nodes"
    fi

    # Execute generate skill
    exec "$(dirname "$0")/generate" "$prompt" "$@"
}

# Execute main with all arguments
main "$@"
