#!/usr/bin/env bash

# Script to generate the arch-toolkit Library Control Flow Diagram
# This script dynamically generates the markdown file with configurable options

set -euo pipefail

# Colors for output (harmonized with Makefile)
COLOR_RESET=$(tput sgr0)
# shellcheck disable=SC2034  # Used in printf statements
COLOR_BOLD=$(tput bold)
COLOR_GREEN=$(tput setaf 2)
COLOR_YELLOW=$(tput setaf 3)
COLOR_BLUE=$(tput setaf 4)

# Configuration variables
OUTPUT_FILE="${OUTPUT_FILE:-../ControlFlow_Diagram.md}"
DIAGRAM_THEME="${DIAGRAM_THEME:-default}"
INCLUDE_DOCS="${INCLUDE_DOCS:-true}"
EXPORT_PNG="${EXPORT_PNG:-false}"
PNG_THEME="${PNG_THEME:-light}"
SRC_DIR="${SRC_DIR:-../../src}"

# Color scheme configuration
COLOR_START="${COLOR_START:-#e1f5ff}"
COLOR_APPRUN="${COLOR_APPRUN:-#fff4e1}"
COLOR_MAINLOOP="${COLOR_MAINLOOP:-#ffe1f5}"
COLOR_SHUTDOWN="${COLOR_SHUTDOWN:-#e1ffe1}"
COLOR_END="${COLOR_END:-#ffe1e1}"

# Get the script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Function to analyze code and extract control flow
analyze_code() {
    local src_path="$1"
    local runtime_file="$src_path/app/runtime.rs"
    
    if [[ ! -f "$runtime_file" ]]; then
        printf "%b❌ Error: Could not find %s%b\n" "$COLOR_YELLOW" "$runtime_file" "$COLOR_RESET" >&2
        return 1
    fi
    
    # Extract main select! branches
    local select_start
    select_start=$(grep -n "select!" "$runtime_file" | head -1 | cut -d: -f1)
    if [[ -z "$select_start" ]]; then
        printf "%b❌ Error: Could not find main select! block in %s%b\n" "$COLOR_YELLOW" "$runtime_file" "$COLOR_RESET" >&2
        return 1
    fi
    
    # Find the select! block and extract branches (look for patterns like "Some(...) = ..._rx.recv()")
    echo "# Extracted from code analysis" >&2
    
    # Extract channel receives in select! block
    local branches
    branches=$(sed -n "${select_start},/^[[:space:]]*}$/p" "$runtime_file" | \
        grep -E "(Some\([^)]+\) = [^=]+\.recv\(\)|Some\(_\) = [^=]+\.recv\(\))" | \
        sed 's/.*Some(\([^)]*\))[[:space:]]*=[[:space:]]*\([^[:space:]]*\)\.recv().*/SELECT_BRANCH:\2:\1/' | \
        sed 's/.*Some(_)[[:space:]]*=[[:space:]]*\([^[:space:]]*\)\.recv().*/SELECT_BRANCH:\1:_/')
    
    # Extract worker spawns
    local workers
    workers=$(grep -n "tokio::spawn" "$runtime_file" | \
        grep -E "(details|deps|files|services|sandbox|pkgb|news|status|tick|query)" | \
        sed 's/.*\(details\|deps\|files\|services\|sandbox\|pkgb\|news\|status\|tick\|query\)[^[:space:]]*/\1_worker/i')
    
    # Return success if we found something
    if [[ -n "$branches" ]] || [[ -n "$workers" ]]; then
        printf "%b✓ Found control flow patterns in code%b\n" "$COLOR_GREEN" "$COLOR_RESET" >&2
        return 0
    else
        printf "%b⚠ Warning: Could not extract control flow patterns, but continuing with basic diagram generation.%b\n" "$COLOR_YELLOW" "$COLOR_RESET" >&2
        return 0  # Still return success, as we can generate a basic diagram
    fi
}

# Function to generate diagram from code analysis
generate_diagram() {
    local src_path="$1"
    local runtime_file="$src_path/app/runtime.rs"
    
    echo '```mermaid'
    echo 'flowchart TD'
    echo '    Start([main.rs: Application Start]) --> ParseArgs[Parse CLI Arguments]'
    echo "    ParseArgs --> CheckFlags{Special Flags?}"
    echo ""
    echo "    CheckFlags -->|--clear-cache| ClearCache[Clear Cache Files & Exit]"
    echo "    CheckFlags -->|--dry-run| SetDryRun[Set Dry Run Flag]"
    echo "    CheckFlags -->|Normal| InitLogging[Initialize Logging]"
    echo ""
    echo "    SetDryRun --> InitLogging"
    echo "    InitLogging --> AppRun[app::run]"
    echo ""
    echo "    AppRun --> SetupTerm[Setup Terminal]"
    echo "    SetupTerm --> InitState[Initialize App State]"
    echo ""
    echo "    InitState --> LoadConfig[Migrate Legacy Configs<br/>Load Settings]"
    echo "    LoadConfig --> LoadLocale[Initialize Locale System]"
    echo "    LoadLocale --> LoadCaches[Load Persisted Caches:<br/>- Details Cache<br/>- Recent Queries<br/>- Install List<br/>- Dependency Cache<br/>- File Cache<br/>- Service Cache<br/>- Sandbox Cache<br/>- News Read URLs<br/>- Official Index]"
    echo ""
    echo "    LoadCaches --> CheckCaches{Caches Valid?}"
    echo "    CheckCaches -->|Missing/Invalid| SetInitFlags[Set Init Flags for<br/>Background Resolution]"
    echo "    CheckCaches -->|Valid| CreateChannels[Create Channels]"
    echo "    SetInitFlags --> CreateChannels"
    echo ""
    echo "    CreateChannels --> SpawnWorkers[Spawn Background Workers:<br/>- Status Worker<br/>- News Worker<br/>- Tick Worker<br/>- Index Update Worker<br/>- Event Reading Thread]"
    echo ""
    echo "    SpawnWorkers --> TriggerResolutions{Init Flags Set?}"
    echo "    TriggerResolutions -->|Yes| SendResolutions[Send Resolution Requests:<br/>- Dependencies<br/>- Files<br/>- Services<br/>- Sandbox]"
    echo "    TriggerResolutions -->|No| SendQuery"
    echo "    SendResolutions --> SendQuery[Send Initial Query]"
    echo ""
    echo "    SendQuery --> MainLoop[Main Event Loop<br/>tokio::select!]"
    echo ""
    echo "    MainLoop --> RenderUI[Render UI Frame]"
    echo "    RenderUI --> SelectEvents{Select on Channels}"
    echo ""
    
    # Extract actual branches from code
    local select_start
    select_start=$(grep -n "select!" "$runtime_file" | head -1 | cut -d: -f1)
    if [[ -n "$select_start" ]]; then
        # Find the end of the select! block (look for closing brace at same indentation)
        local select_end
        select_end=$(awk -v start="$select_start" '
            NR >= start {
                if (match($0, /^[[:space:]]*select![[:space:]]*\{/)) { depth=1; next }
                if (match($0, /\{/)) { depth++ }
                if (match($0, /\}/)) { depth--; if (depth == 0) { print NR; exit } }
            }
        ' "$runtime_file")
        
        if [[ -n "$select_end" ]]; then
            # Extract branches and create diagram nodes
            # Use sed/grep to extract channel names, then process them
            sed -n "${select_start},${select_end}p" "$runtime_file" | \
            grep -E "Some\([^)]+\) = [^=]+\.recv\(\)" | \
            sed -E 's/.*Some\(([^)]+)\)[[:space:]]*=[[:space:]]*([^[:space:]]+)_rx\.recv\(\).*/\2:\1/' | \
            while IFS=: read -r channel _; do
                if [[ -n "$channel" ]]; then
                    local handler_name
                    handler_name=$(echo "$channel" | sed 's/_res$//' | sed 's/_rx$//' | sed 's/_notify$//' | sed 's/_//g')
                    # Capitalize first letter
                    handler_name="${handler_name^}"
                    
                    # Determine handler type based on channel name
                    if [[ "$channel" =~ _res$ ]]; then
                        echo "    SelectEvents -->|${channel^} Result| Handle${handler_name}[Handle ${handler_name^} Result]"
                    elif [[ "$channel" =~ _notify$ ]]; then
                        echo "    SelectEvents -->|${channel^} Notify| Update${handler_name}[Update ${handler_name^} State]"
                    else
                        echo "    SelectEvents -->|${channel^}| Handle${handler_name}[Handle ${handler_name^}]"
                    fi
                fi
            done
        fi
    fi
    
    # Add standard branches (fallback if extraction fails or for common patterns)
    # These are always included as they're core to the application flow
    echo "    SelectEvents -->|Event Received| HandleEvent[events::handle_event]"
    
    # Try to extract actual branches, but include common ones as fallback
    if [[ -n "$select_start" ]] && [[ -n "$select_end" ]]; then
        local extracted
        extracted=$(sed -n "${select_start},${select_end}p" "$runtime_file" | \
            grep -E "Some\([^)]+\) = [^=]+\.recv\(\)" | \
            sed -E 's/.*Some\(([^)]+)\)[[:space:]]*=[[:space:]]*([^[:space:]]+)_rx\.recv\(\).*/\2:\1/' | \
            head -1)
        # extracted is checked but not used further - this is intentional for future use
        # shellcheck disable=SC2034
        if [[ -n "$extracted" ]]; then
            : # Branch extraction successful
        fi
    fi
    
    # Always include these core branches
    echo "    SelectEvents -->|Index Notify| UpdateIndex[Update Index State]"
    echo "    SelectEvents -->|Search Results| HandleResults[Handle Search Results]"
    echo "    SelectEvents -->|Details Update| HandleDetails[Handle Details Update]"
    echo "    SelectEvents -->|Preview Update| HandlePreview[Handle Preview]"
    echo "    SelectEvents -->|Add to Install| HandleAdd[Batch Add to Install List<br/>Trigger Resolutions]"
    echo "    SelectEvents -->|Deps Result| HandleDeps[Handle Dependency Result]"
    echo "    SelectEvents -->|Files Result| HandleFiles[Handle File Result]"
    echo "    SelectEvents -->|Services Result| HandleServices[Handle Service Result]"
    echo "    SelectEvents -->|Sandbox Result| HandleSandbox[Handle Sandbox Result]"
    echo "    SelectEvents -->|PKGBUILD Result| HandlePKGBUILD[Handle PKGBUILD Result]"
    echo "    SelectEvents -->|Summary Result| HandleSummary[Handle Summary Result]"
    echo "    SelectEvents -->|Network Error| ShowAlert[Show Alert Modal]"
    echo "    SelectEvents -->|Tick Event| HandleTick[Handle Tick Event]"
    echo "    SelectEvents -->|News Update| HandleNews[Handle News Update]"
    echo "    SelectEvents -->|Status Update| HandleStatus[Handle Status Update]"
    echo ""
    echo "    HandleEvent --> CheckModal{Modal Active?}"
    echo "    CheckModal -->|Yes| HandleModal[Handle Modal Events]"
    echo "    CheckModal -->|No| CheckGlobal{Global Shortcut?}"
    echo ""
    echo "    HandleModal --> CheckExit{Exit Requested?}"
    echo "    CheckGlobal -->|Yes| HandleGlobal[Handle Global Shortcuts]"
    echo "    CheckGlobal -->|No| HandlePane[Handle Pane-Specific Events:<br/>- Search Pane<br/>- Recent Pane<br/>- Install Pane]"
    echo ""
    echo "    HandleGlobal --> CheckExit"
    echo "    HandlePane --> CheckExit"
    echo "    HandleModal --> CheckExit"
    echo ""
    echo "    CheckExit -->|Yes| Shutdown"
    echo "    CheckExit -->|No| RenderUI"
    echo ""
    echo "    UpdateIndex --> RenderUI"
    echo "    HandleResults --> RenderUI"
    echo "    HandleDetails --> RenderUI"
    echo "    HandlePreview --> RenderUI"
    echo "    HandleAdd --> RenderUI"
    echo "    HandleDeps --> RenderUI"
    echo "    HandleFiles --> RenderUI"
    echo "    HandleServices --> RenderUI"
    echo "    HandleSandbox --> RenderUI"
    echo "    HandlePKGBUILD --> RenderUI"
    echo "    HandleSummary --> RenderUI"
    echo "    ShowAlert --> RenderUI"
    echo "    HandleNews --> RenderUI"
    echo "    HandleStatus --> RenderUI"
    echo ""
    echo "    HandleTick --> TickTasks[Periodic Tasks:<br/>- Flush Caches<br/>- Save Recent<br/>- Preflight Resolution<br/>- PKGBUILD Debounce<br/>- Poll Installed Cache<br/>- Ring Prefetch<br/>- Auto-close Menus<br/>- Expire Toasts]"
    echo "    TickTasks --> RenderUI"
    echo ""
    echo "    Shutdown[Shutdown Sequence] --> ResetFlags[Reset Resolution Flags]"
    echo "    ResetFlags --> SignalEventThread[Signal Event Thread to Exit]"
    echo "    SignalEventThread --> FlushAll[Flush All Caches:<br/>- Details Cache<br/>- Recent Queries<br/>- Install List<br/>- News Read URLs<br/>- Dependency Cache<br/>- File Cache<br/>- Service Cache<br/>- Sandbox Cache]"
    echo "    FlushAll --> RestoreTerm[Restore Terminal]"
    echo "    RestoreTerm --> End([Exit])"
    echo ""
    echo "    ClearCache --> End"
    echo ""
    echo "    style Start fill:$COLOR_START"
    echo "    style AppRun fill:$COLOR_APPRUN"
    echo "    style MainLoop fill:$COLOR_MAINLOOP"
    echo "    style Shutdown fill:$COLOR_SHUTDOWN"
    echo "    style End fill:$COLOR_END"
    echo '```'
}


# Function to generate documentation sections
generate_documentation() {
    cat <<'DOC_EOF'

## Key Components

### 1. Initialization Phase
- **CLI Argument Parsing**: Handles command-line flags (--dry-run, --clear-cache, etc.)
- **Logging Setup**: Initializes tracing logger to file
- **Terminal Setup**: Configures terminal for TUI mode
- **State Initialization**: Loads settings, caches, locale system
- **Channel Creation**: Sets up async communication channels
- **Worker Spawning**: Launches background workers for async operations

### 2. Main Event Loop
The application uses `tokio::select!` to concurrently handle multiple async channels:
- **User Input**: Keyboard and mouse events
- **Search Results**: Package search results from AUR/official repos
- **Details Updates**: Package information updates
- **Analysis Results**: Dependency, file, service, and sandbox analysis
- **PKGBUILD Content**: Package build file content
- **Preflight Summary**: Installation preflight analysis results
- **News/Status**: Arch Linux news and status updates
- **Tick Events**: Periodic background tasks

### 3. Event Handling
Events are processed in priority order:
1. **Modal Interactions**: Active modal dialogs (handled first)
2. **Global Shortcuts**: Application-wide shortcuts (help, exit, theme reload)
3. **Pane-Specific Events**: Search, Recent, and Install pane interactions

### 4. Background Workers
Asynchronous workers handle:
- **Search Worker**: AUR and official repository package search
- **Details Worker**: Package information retrieval
- **Dependency Worker**: Dependency resolution and analysis
- **File Worker**: File system impact analysis
- **Service Worker**: Systemd service impact analysis
- **Sandbox Worker**: AUR package sandbox analysis
- **News Worker**: Arch Linux news fetching
- **Status Worker**: Arch status page monitoring
- **Index Worker**: Official package index updates

### 5. Tick Handler (Periodic Tasks)
The tick handler performs periodic maintenance:
- **Cache Persistence**: Debounced writes of dirty caches
- **Preflight Resolution**: Processes queued preflight analysis requests
- **PKGBUILD Debouncing**: Manages PKGBUILD reload requests
- **Installed Cache Polling**: Refreshes installed package cache after installs/removals
- **Ring Prefetch**: Prefetches details for packages around selection
- **UI State Cleanup**: Auto-closes menus and expires toast messages

### 6. Shutdown Sequence
Graceful shutdown process:
- Reset all resolution flags
- Signal background threads to exit
- Flush all pending cache writes
- Restore terminal to original state

## Architecture Notes

- **Async Architecture**: Uses Tokio for async runtime with channels for communication
- **Event-Driven**: Main loop responds to events from multiple sources
- **Background Processing**: Heavy I/O operations run in background workers
- **State Management**: Centralized `AppState` holds all application state
- **Cache Strategy**: Multiple caches with signature-based validation
- **Debouncing**: Used for cache writes and PKGBUILD reloads to reduce I/O

## Converting to Image

To convert this Mermaid diagram to a PNG image, you can use:

1. **Mermaid CLI**: `mmdc -i ControlFlow_Diagram.md -o ControlFlow_Diagram.png`
2. **Online Tools**: Paste the mermaid code block into https://mermaid.live/
3. **VS Code Extension**: Use the "Markdown Preview Mermaid Support" extension
4. **GitHub/GitLab**: The diagram will render automatically in markdown files
DOC_EOF
}

# Main function to generate the markdown file
generate_markdown() {
    {
        echo "# arch-toolkit Library Control Flow Diagram"
        echo ""
        echo "This diagram shows the complete control flow of the arch-toolkit library operations."
        echo ""
        echo "> **Note**: This diagram was generated by analyzing the codebase. It reflects the actual control flow in the source code."
        echo ""
        
        # Resolve source directory path
        local src_path
        if [[ "$SRC_DIR" == /* ]]; then
            src_path="$SRC_DIR"
        elif [[ "$SRC_DIR" == ../* ]] || [[ "$SRC_DIR" == ./* ]]; then
            src_path="$(cd "$SCRIPT_DIR" && cd "$(dirname "$SRC_DIR")" && pwd)/$(basename "$SRC_DIR")"
        else
            src_path="$(cd "$SCRIPT_DIR/../.." && pwd)/$SRC_DIR"
        fi
        
        # Always analyze code and generate diagram
        if analyze_code "$src_path"; then
            generate_diagram "$src_path" | sed \
                -e "s/COLOR_START/$COLOR_START/g" \
                -e "s/COLOR_APPRUN/$COLOR_APPRUN/g" \
                -e "s/COLOR_MAINLOOP/$COLOR_MAINLOOP/g" \
                -e "s/COLOR_SHUTDOWN/$COLOR_SHUTDOWN/g" \
                -e "s/COLOR_END/$COLOR_END/g"
        else
            printf "%b❌ Error: Could not analyze codebase. Please check that %s exists and contains the runtime.rs file.%b\n" "$COLOR_YELLOW" "$src_path" "$COLOR_RESET" >&2
            exit 1
        fi
        
        # Add documentation if requested
        if [[ "$INCLUDE_DOCS" == "true" ]]; then
            generate_documentation
        fi
    } > "$OUTPUT_FILE_ABS"
    
    printf "%b✓ Generated markdown file: %s%b\n" "$COLOR_GREEN" "$OUTPUT_FILE_ABS" "$COLOR_RESET"
}

# Function to export to PNG if mermaid-cli is available
export_to_png() {
    if ! command -v mmdc &> /dev/null; then
        printf "%b⚠ Warning: mermaid-cli (mmdc) not found. Skipping PNG export.%b\n" "$COLOR_YELLOW" "$COLOR_RESET"
        echo "  Install with: npm install -g @mermaid-js/mermaid-cli"
        return 1
    fi
    
    local png_output="${OUTPUT_FILE_ABS%.md}.png"
    local theme_flag=""
    
    if [[ "$PNG_THEME" == "dark" ]]; then
        theme_flag="--theme dark"
    elif [[ "$PNG_THEME" == "light" ]]; then
        theme_flag="--theme light"
    fi
    
    printf "%bExporting to PNG (theme: %s)...%b\n" "$COLOR_BLUE" "$PNG_THEME" "$COLOR_RESET"
    # shellcheck disable=SC2086  # theme_flag may be empty or contain multiple words
    mmdc -i "$OUTPUT_FILE_ABS" -o "$png_output" $theme_flag 2>/dev/null || {
        printf "%b⚠ Warning: PNG export failed. Continuing...%b\n" "$COLOR_YELLOW" "$COLOR_RESET"
        return 1
    }
    
    printf "%b✓ Generated PNG file: %s%b\n" "$COLOR_GREEN" "$png_output" "$COLOR_RESET"
}

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --output|-o)
            OUTPUT_FILE="$2"
            shift 2
            ;;
        --theme)
            DIAGRAM_THEME="$2"
            shift 2
            ;;
        --no-docs)
            INCLUDE_DOCS="false"
            shift
            ;;
        --src-dir)
            SRC_DIR="$2"
            shift 2
            ;;
        --export-png)
            EXPORT_PNG="true"
            shift
            ;;
        --png-theme)
            PNG_THEME="$2"
            EXPORT_PNG="true"
            shift 2
            ;;
        --color-start)
            COLOR_START="$2"
            shift 2
            ;;
        --color-apprun)
            COLOR_APPRUN="$2"
            shift 2
            ;;
        --color-mainloop)
            COLOR_MAINLOOP="$2"
            shift 2
            ;;
        --color-shutdown)
            COLOR_SHUTDOWN="$2"
            shift 2
            ;;
        --color-end)
            COLOR_END="$2"
            shift 2
            ;;
        --help|-h)
            cat <<EOF
Usage: $0 [OPTIONS]

Generate the arch-toolkit Library Control Flow Diagram markdown file.

OPTIONS:
    -o, --output FILE          Output file path (default: ../ControlFlow_Diagram.md)
    --theme THEME              Diagram theme (default: default)
    --no-docs                  Exclude documentation sections
    --src-dir DIR              Source directory path (default: ../../src)
    --export-png               Export diagram to PNG after generation
    --png-theme THEME          PNG theme: light, dark, or default (implies --export-png)
    --color-start COLOR        Color for Start node (default: #e1f5ff)
    --color-apprun COLOR       Color for AppRun node (default: #fff4e1)
    --color-mainloop COLOR     Color for MainLoop node (default: #ffe1f5)
    --color-shutdown COLOR     Color for Shutdown node (default: #e1ffe1)
    --color-end COLOR          Color for End node (default: #ffe1e1)
    -h, --help                 Show this help message

EXAMPLES:
    # Generate with default settings
    $0

    # Generate with custom colors
    $0 --color-start "#ff0000" --color-end "#00ff00"

    # Generate and export to PNG
    $0 --export-png --png-theme dark

    # Generate without documentation
    $0 --no-docs

    # Use custom source directory
    $0 --src-dir /path/to/src

ENVIRONMENT VARIABLES:
    All options can also be set via environment variables:
    OUTPUT_FILE, DIAGRAM_THEME, INCLUDE_DOCS, EXPORT_PNG, PNG_THEME, SRC_DIR,
    COLOR_START, COLOR_APPRUN, COLOR_MAINLOOP, COLOR_SHUTDOWN, COLOR_END
EOF
            exit 0
            ;;
        *)
            printf "%bUnknown option: %s%b\n" "$COLOR_YELLOW" "$1" "$COLOR_RESET" >&2
            echo "Use --help for usage information" >&2
            exit 1
            ;;
    esac
done

# Main execution
cd "$SCRIPT_DIR"

# Resolve output file path (after changing to script directory)
if [[ "$OUTPUT_FILE" == /* ]]; then
    # Absolute path
    OUTPUT_FILE_ABS="$OUTPUT_FILE"
elif [[ "$OUTPUT_FILE" == ../* ]]; then
    # Relative to script directory (go up from scripts/)
    OUTPUT_FILE_ABS="$(cd "$(dirname "$OUTPUT_FILE")" && pwd)/$(basename "$OUTPUT_FILE")"
elif [[ "$OUTPUT_FILE" == ./* ]]; then
    # Relative to script directory
    OUTPUT_FILE_ABS="$(pwd)/${OUTPUT_FILE#./}"
else
    # If it doesn't start with ../ or ./, try to resolve it
    # First try relative to script directory, then relative to project root
    if [[ -f "$(pwd)/$OUTPUT_FILE" ]] || [[ "$OUTPUT_FILE" != */* ]]; then
        # File in script directory or just a filename
        OUTPUT_FILE_ABS="$(pwd)/$OUTPUT_FILE"
    else
        # Assume relative to project root (parent of dev/)
        PROJECT_ROOT="$(cd ../.. && pwd)"
        OUTPUT_FILE_ABS="$PROJECT_ROOT/$OUTPUT_FILE"
    fi
fi

generate_markdown

if [[ "$EXPORT_PNG" == "true" ]]; then
    export_to_png
fi

printf "%b✓ Done!%b\n" "$COLOR_GREEN" "$COLOR_RESET"