#!/bin/bash

# Fission Docstring Updater
# Parses and updates embedded fission configuration in Python function docstrings

set -euo pipefail

usage() {
    echo "Usage: $0 <file-path> [function-name] [--set \"<json>\"] [--get] [--help]"
    echo ""
    echo "Arguments:"
    echo "  file-path:     Path to the Python file containing the function"
    echo "  function-name: Optional specific function name to target (if not provided, processes all functions with fission configuration)"
    echo ""
    echo "Options:"
    echo "  --set <json>:  Set the fission configuration to the provided JSON string"
    echo "  --get:         Get/display the current fission configuration (default action)"
    echo "  --help:        Show this help message"
    echo ""
    echo "Examples:"
    echo "  $0 ./src/my_function.py main --get"
    echo "  $0 ./src/my_function.py main --set '{\"name\": \"updated-function\"}'"
    echo "  $0 ./src/functions.py --get"
    exit 1
}

# Check dependencies
if ! command -v python3 &> /dev/null; then
    echo "Error: python3 is required but not found"
    exit 1
fi

# Parse arguments
if [[ $# -lt 1 ]]; then
    usage
fi

FILE_PATH="$1"
shift

FUNCTION_NAME=""
ACTION="get"
SET_VALUE=""

while [[ $# -gt 0 ]]; do
    case $1 in
        --set)
            SET_VALUE="$2"
            ACTION="set"
            shift 2
            ;;
        --get)
            ACTION="get"
            shift
            ;;
        --help)
            usage
            ;;
        *)
            if [[ -z "$FUNCTION_NAME" ]]; then
                FUNCTION_NAME="$1"
            else
                echo "Error: Unexpected argument '$1'"
                usage
            fi
            shift
            ;;
    esac
done

# Validate file exists
if [[ ! -f "$FILE_PATH" ]]; then
    echo "Error: File '$FILE_PATH' does not exist"
    exit 1
fi

# Validate JSON if setting
if [[ "$ACTION" == "set" && -z "$SET_VALUE" ]]; then
    echo "Error: --set requires a JSON value"
    exit 1
fi

if [[ "$ACTION" == "set" ]]; then
    # Validate JSON format
    if ! echo "$SET_VALUE" | python3 -m json.tool >/dev/null 2>&1; then
        echo "Error: Invalid JSON provided for --set"
        exit 1
    fi
fi

# Python script to handle the docstring parsing and updating
PYTHON_SCRIPT=$(cat << 'EOF'
import re
import sys
import json
import os
from pathlib import Path

def extract_functions_with_fission(content):
    """Extract all functions that have fission configuration in their docstrings."""
    # Pattern to match Python functions with docstrings containing fission configuration
    # This looks for def function_name(): followed by a docstring that contains ```fission
    pattern = r'(\s*def\s+(\w+)\s*\([^)]*\):\s*(?:\n\s*)?\"\"\"[\s\S]*?```fission[\s\S]*?```[\s\S]*?\"\"\"[\s\S]*?)(?=\n\s*def|\n\s*class|\Z)'

    matches = re.finditer(pattern, content, re.MULTILINE)
    functions = []

    for match in matches:
        full_match = match.group(1)
        # Extract function name from the match
        func_name_match = re.search(r'def\s+(\w+)\s*\(', full_match)
        if func_name_match:
            func_name = func_name_match.group(1)
            functions.append({
                'name': func_name,
                'full_text': full_match,
                'start_pos': match.start(),
                'end_pos': match.end()
            })

    return functions

def extract_fission_config(docstring):
    """Extract fission configuration from a docstring."""
    # Look for ```fission ... ``` blocks
    pattern = r'```fission\s*([\s\S]*?)\s*```'
    match = re.search(pattern, docstring)
    if match:
        config_text = match.group(1).strip()
        try:
            return json.loads(config_text)
        except json.JSONDecodeError as e:
            return None
    return None

def replace_fission_config_in_docstring(docstring, new_config):
    """Replace fission configuration in a docstring with new config."""
    # Format the new config as JSON with indentation
    formatted_config = json.dumps(new_config, indent=4)
    # Replace the ```fission ... ``` block
    pattern = r'(```fission\s*)[\s\S]*?(\s*```)'
    replacement = r'\1' + formatted_config + r'\2'
    return re.sub(pattern, replacement, docstring, flags=re.DOTALL)

def process_file(file_path, target_function=None, action='get', set_value=None):
    """Process the Python file to get or set fission configuration."""
    try:
        with open(file_path, 'r') as f:
            content = f.read()
    except IOError as e:
        print(f"Error: Cannot read file '{file_path}': {e}", file=sys.stderr)
        sys.exit(1)

    functions = extract_functions_with_fission(content)

    if not functions:
        print("No functions with fission configuration found in file.", file=sys.stderr)
        if action == 'get':
            sys.exit(0)
        else:
            sys.exit(1)

    # Filter by function name if specified
    if target_function:
        functions = [f for f in functions if f['name'] == target_function]
        if not functions:
            print(f"Error: Function '{target_function}' with fission configuration not found.", file=sys.stderr)
            sys.exit(1)

    if action == 'get':
        # Display current configuration for each function
        for func in functions:
            # Extract docstring from the function text
            docstring_match = re.search(r'\"\"\"[\s\S]*?\"\"\"', func['full_text'])
            if docstring_match:
                docstring = docstring_match.group(0)
                config = extract_fission_config(docstring)
                if config is not None:
                    if len(functions) == 1:
                        print(json.dumps(config, indent=2))
                    else:
                        print(f"Function '{func['name']}':")
                        print(json.dumps(config, indent=2))
                        print()
                else:
                    if len(functions) == 1:
                        print("No fission configuration found in function docstring.", file=sys.stderr)
                    else:
                        print(f"Function '{func['name']}': No fission configuration found in docstring.", file=sys.stderr)
            else:
                if len(functions) == 1:
                    print("Could not extract docstring from function.", file=sys.stderr)
                else:
                    print(f"Function '{func['name']}': Could not extract docstring.", file=sys.stderr)

    elif action == 'set':
        if set_value is None:
            print("Error: No value provided for --set", file=sys.stderr)
            sys.exit(1)

        try:
            new_config = json.loads(set_value)
        except json.JSONDecodeError as e:
            print(f"Error: Invalid JSON provided for --set: {e}", file=sys.stderr)
            sys.exit(1)

        # Update each function
        updated_content = content
        offset = 0  # Track position changes due to replacements

        for func in functions:
            # Extract docstring from the function text
            docstring_match = re.search(r'\"\"\"[\s\S]*?\"\"\"', func['full_text'])
            if docstring_match:
                docstring = docstring_match.group(0)
                # Check if fission configuration exists
                if extract_fission_config(docstring) is not None:
                    # Replace fission configuration in docstring
                    new_docstring = replace_fission_config_in_docstring(docstring, new_config)
                    # Replace the docstring in the function text
                    new_func_text = func['full_text'].replace(docstring, new_docstring, 1)
                    # Replace in the overall content (adjusting for previous changes)
                    start_pos = func['start_pos'] + offset
                    end_pos = func['end_pos'] + offset
                    # Update content with the change
                    before = updated_content[:start_pos]
                    after = updated_content[end_pos:]
                    updated_content = before + new_func_text + after
                    # Update offset for next replacements
                    offset += len(new_func_text) - len(func['full_text'])
                else:
                    print(f"Warning: No fission configuration found in function '{func['name']}' to update.", file=sys.stderr)
            else:
                print(f"Warning: Could not extract docstring from function '{func['name']}'.", file=sys.stderr)

        # Write back to file
        try:
            with open(file_path, 'w') as f:
                f.write(updated_content)
            if len(functions) == 1:
                print(f"Updated fission configuration in function '{functions[0]['name']}'.")
            else:
                print(f"Updated fission configuration in {len(functions)} function(s).")
        except IOError as e:
            print(f"Error: Cannot write to file '{file_path}': {e}", file=sys.stderr)
            sys.exit(1)

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print("Usage: fission-docstring-updater <file-path> [function-name] [--set \"<json>\"] [--get]", file=sys.stderr)
        sys.exit(1)

    file_path = sys.argv[1]
    target_function = sys.argv[2] if len(sys.argv) > 2 and not sys.argv[2].startswith('--') else None

    # Parse arguments
    action = 'get'
    set_value = None

    i = 3 if target_function else 2
    while i < len(sys.argv):
        if sys.argv[i] == '--set' and i + 1 < len(sys.argv):
            action = 'set'
            set_value = sys.argv[i + 1]
            i += 2
        elif sys.argv[i] == '--get':
            action = 'get'
            i += 1
        else:
            print(f"Error: Unknown argument '{sys.argv[i]}'", file=sys.stderr)
            sys.exit(1)

    process_file(file_path, target_function, action, set_value)
EOF
)

# Build arguments for Python script
PYTHON_ARGS=("$FILE_PATH")
if [[ -n "$FUNCTION_NAME" ]]; then
    PYTHON_ARGS+=("$FUNCTION_NAME")
fi

if [[ "$ACTION" == "set" ]]; then
    PYTHON_ARGS+=("--set" "$SET_VALUE")
else
    PYTHON_ARGS+=("--get")
fi

# Execute the Python script
python3 -c "$PYTHON_SCRIPT" "${PYTHON_ARGS[@]}"