#!/bin/bash

function help() {
    echo
    echo "$1 <command> [options]"
    echo
    echo "commands:"
    echo

    CMD_WIDTH=23
    OPT_WIDTH=22
    DESC_WIDTH=70
    COLUMN="| %-${CMD_WIDTH}s | %-${OPT_WIDTH}s | %-${DESC_WIDTH}s |\n"

    # Print table header
    printf "$COLUMN" "Command" "Option" "Description"
    echo "|$(printf '%*s' $((CMD_WIDTH + 2)) '' | tr ' ' '-')|$(printf '%*s' $((OPT_WIDTH + 2)) '' | tr ' ' '-')|$(printf '%*s' $((DESC_WIDTH + 2)) '' | tr ' ' '-')|"

    # Print table rows
    printf "$COLUMN" "generate-vault-password" "[vault_id] [password]" "Generate a password file for Ansible vault."
    printf "$COLUMN" "convert-vault-file" "[vault_id] [path]" "Re-encrypt vault file as single strings."
    printf "$COLUMN" "show-vault-file" "[vault_id] [path]" "Show content of encrypted strings or fully encrypted vault file."
    printf "$COLUMN" "encrypt-string" "[vault_id] [key:value] -q" "Generate encrypted string with Ansible Vault."
    printf "$COLUMN" "install" "" "Install Ansible and dependencies with uv."
    printf "$COLUMN" "install-exo" "" "Install Exoscale CLI."
    printf "$COLUMN" "upgrade" "" "Upgrade Ansible and dependencies with uv."
    printf "$COLUMN" "update-with-llm" "[path][prompt]" "Update files with LLM prompt."
    printf "$COLUMN" "lint" "" "Lint Ansible roles and playbooks."
    printf "$COLUMN" "build-galaxy" "" "Create tarball of collection."
    printf "$COLUMN" "publish-galaxy" "" "Publish tarball to Ansible Galaxy."
    printf "$COLUMN" "install-node" "" "Install build dependencies."
    printf "$COLUMN" "dev-node" "" "Start vuepress development server."
    printf "$COLUMN" "build-node" "" "Create vuepress build."
    printf "$COLUMN" "serve-node-build" "" "Serve vuepress build."
    printf "$COLUMN" "list-vault" "" "Lists vault identities."
    printf "$COLUMN" "list-hosts" "[path]" "Lists all hosts of an inventory."
    printf "$COLUMN" "version" "" "Show tool versions."
    printf "$COLUMN" "source" "" "Source the Python virtual env."
    printf "$COLUMN" "test-localhost" "[clean]" "Deploy to localhost. Pass option 'clean' to clean up."
    printf "$COLUMN" "test-scripts" "[clean]" "Install and test scripts."
    printf "$COLUMN" "show-passwordless-sudo" "[username]" "Give instructions to setup passwordless sudo."

    echo
}

# Set env vars
GIT_BRANCH=$(git symbolic-ref --short -q HEAD)
LLM_MODEL=${LLM_MODEL:="deepseek-coder"}

if [[ -a ".env" ]]; then
    export $(cat .env | sed 's/^#.*//g' | xargs)
fi

# Ansible Vault commands

function list-vault() {
    grep 'vault_identity_list' ansible.cfg | sed 's/^vault_identity_list = //' | tr ',' '\n' | sed '/^$/d'
}

function generate-vault-password() {
    if test -z "$1"; then echo "\$1 is empty"; exit; fi
    if test -z "$2"; then echo "\$1 is empty"; exit; fi

    export VAULT_ID=$1
    export VAULTPASSWORD=$2
    echo "$VAULTPASSWORD" > ".vault_pass_$VAULT_ID"
    chmod 600 ".vault_pass_$VAULT_ID"
    echo "Password file ".vault_pass_$VAULT_ID" created."
}

function encrypt-string() {
    if test -z "$1"; then echo "\$1 is empty"; exit; fi
    if test -z "$2"; then echo "\$2 is empty"; exit; fi

    # Split multiline
    STRING_KEY=$(echo "$2" | cut -d ':' -f 1)
    STRING_VALUE=$(echo "$2" | cut -d ':' -f 2 | sed 's/^ *//')

    if [ "$3" != "-q" ]; then 
        echo "Found key: $STRING_KEY"
        echo "Encrypting value: $STRING_VALUE"
    fi
    ansible-vault encrypt_string --encrypt-vault-id="$1"  --name "$STRING_KEY" "$STRING_VALUE"
}

function convert-vault-file() {
    if test -z "$1"; then echo "\$1 is empty"; exit; fi
    if test -z "$2"; then echo "\$2 is empty"; exit; fi
    activate-venv
    
    VAULT_ID=$1
    FILE_PATH=$2
    TEMP_FILE=$(mktemp)

    # Decrypt the vault file
    ansible-vault decrypt --output="$FILE_PATH" "$FILE_PATH"

    # Process each key-value pair
    yq eval '.' "$FILE_PATH" | while IFS= read -r LINE; do
        KEY=$(echo "$LINE" | cut -d':' -f1)
        VALUE=$(echo "$LINE" | cut -d':' -f2-)
        VALUE=$(echo "$VALUE" | sed 's/^[[:space:]"]*//;s/[[:space:]"]*$//')
        
        # Re-encrypt the key-value pair
        ENCYRPTED=$(ansible-vault encrypt_string --encrypt-vault-id="$VAULT_ID" --name "$KEY" "$VALUE" 2>/dev/null)
        echo "$ENCYRPTED" >> "$TEMP_FILE"
    done

    # Move temp file to vault file
    mv "$TEMP_FILE" "$FILE_PATH"
    echo "Converted file saved as $FILE_PATH"
}

show-vault-file() {
    if test -z "$1"; then echo "\$1 is empty"; exit 1; fi
    if test -z "$2"; then echo "\$2 is empty"; exit 1; fi

    VAULT_ID="$1"
    FILE_PATH="$2"
    TEMP_FILE=$(mktemp)
    TEMP_PART_FILE=$(mktemp)

    if head -n 1 "$FILE_PATH" | grep -q "^\$ANSIBLE_VAULT;1.2;AES256;$VAULT_ID"; then
        ansible-vault decrypt --output="$TEMP_FILE" "$FILE_PATH"
    else
        WRITE_FINISHED=false
        while IFS= read -r LINE; do
            if [[ "$LINE" =~ "!vault" ]]; then
                if [ "$WRITE_FINISHED" ] && [ -s "$TEMP_PART_FILE" ]; then
                    ansible-vault decrypt "$TEMP_PART_FILE"
                    KEY=$(echo "$LINE" | cut -d':' -f1)
                    VALUE=$(cat "$TEMP_PART_FILE")
                    echo "$KEY: $VALUE" >> "$TEMP_FILE"
                fi
                : > "$TEMP_PART_FILE"
                WRITE_FINISHED=true
            else
                if [ "$WRITE_FINISHED" ]; then
                    echo "$LINE" | sed 's/^[[:space:]"]*//;s/[[:space:]"]*$//' >> "$TEMP_PART_FILE"
                fi
            fi
        done < "$FILE_PATH"
    fi

    batcat --paging=always "$TEMP_FILE"
    rm -f "$TEMP_FILE"
    rm -f "$TEMP_PART_FILE"
}

# Environment commands

function init-venv() {
    if [ ! -d "venv$GIT_BRANCH" ]; then
        echo "Init venv$GIT_BRANCH with $(uv --version)."
        uv venv "venv$GIT_BRANCH"
    fi
}

function activate-venv() {
    echo "Source virtualenv venv$GIT_BRANCH."
    source "venv$GIT_BRANCH/bin/activate"
}

function install() {
    echo "Setup Ansible log file"
    sudo touch /var/log/ansible.log
    sudo chown $USER: /var/log/ansible.log

    init-venv
    activate-venv

    echo "Install Ansible and Python dependencies"
    uv pip install -r requirements.txt

    ansible-galaxy collection install community.general
}

function install-exo() {
    curl -fsSL https://raw.githubusercontent.com/exoscale/cli/master/install-latest.sh | sh
}

function version() {
    activate-venv
    uv --version
    ansible --version
}

function show-passwordless-sudo() {
    USER=${1:-"$USERNAME"}
    echo "echo \"$USER ALL=(ALL) NOPASSWD: ALL\" >> /etc/sudoers"
}

# Ansible commands

function upgrade() {
    echo "Update Ansible"
    uv pip install --upgrade ansible
}

# Ansible Invenotry commands

function list-hosts() {
    if test -z "$1"; then echo "\$1 is empty"; exit; fi

    activate-venv
    ansible -i $1 all --list-hosts
}

# Ansible Galaxy commands

function build-galaxy() {
    activate-venv
    VERSION=$(yq eval '.version' galaxy.yml)
    NAMESPACE=$(yq eval '.namespace' galaxy.yml)
    NAME=$(yq eval '.name' galaxy.yml)
    
    ansible-galaxy collection build
    mv $NAMESPACE-$NAME-$VERSION.tar.gz tmp/
}

function publish-galaxy() {
    activate-venv
    VERSION=$(yq eval '.version' galaxy.yml)
    NAMESPACE=$(yq eval '.namespace' galaxy.yml)
    NAME=$(yq eval '.name' galaxy.yml)

    ansible-galaxy collection publish tmp/$NAMESPACE-$NAME-$VERSION.tar.gz --token "$GALAXY_API_TOKEN"
}

# Test commands

function test-localhost() {
    activate-venv
    version
    if [ "$1" = "clean" ]; then
        echo "Cleanup localhost"
        ansible-playbook -i test plays/clean.yml -t "odoo,odoo_data,odoo_volume,\
            postgres,postgres_data,postgres_volume,\
            meilisearch,meilisearch_data,meilisearch_volume,\
            pgadmin,pgadmin_data,pgadmin_volume,\
            n8n,n8n_data,n8n_volume,\
            nginx,nginx_data,nginx_volume,\
            metabase,metabase_data,metabase_volume,\
            prometheus,prometheus_data,prometheus_volume,\
            grafana,grafana_data,grafana_volume,\
            loki,loki_data,loki_volume,\
            promtail,promtail_data,promatil_volume"
    else
        echo "Setup localhost"
        ansible-playbook -i test plays/localhost.yml --skip-tags check

        echo "Curl the services and check HTTP code."
        [[ $(curl -o /dev/null -s -w "%{http_code}" -L "http://n8n.local:8080") -eq 200 ]]
        [[ $(curl -o /dev/null -s -w "%{http_code}" -L "http://odoo.local:8080") -eq 200 ]]
        [[ $(curl -o /dev/null -s -w "%{http_code}" -L "http://odoo.local:8080/website/info") -eq 403 ]]
        [[ $(curl -o /dev/null -s -w "%{http_code}" -L "http://metabase.local:8080") -eq 200 ]]
        [[ $(curl -o /dev/null -s -w "%{http_code}" -L "http://grafana.local:8080") -eq 200 ]]
    fi
}

function test-scripts() {
    set -e
    activate-venv

    if [ "$1" = "clean" ]; then
        echo "Cleanup localhost"
        ansible localhost -m ansible.builtin.cron -a 'name="test" state="absent"'
        ansible-playbook -i test plays/clean.yml -t "odoo,odoo_data,odoo_volume,\
            postgres,postgres_data,postgres_volume"
    else
        echo -e "\033[38;5;214mTEST\033[0m: Test restic scripts"
        sudo ./roles/cron/files/install

        # Add cron job
        ansible localhost -m ansible.builtin.cron -a 'name="test" minute="0" hour="*/6" job="echo \"test\""'

        # List cron job and compare output
        [[ "$(cron-job-list)" == "0 */6 * * * test" ]]

        echo -e "\033[38;5;214mTEST\033[0m: Test postgres scripts"
        sudo ./roles/postgres/files/install

        # Start postgres container
        ansible-playbook -i test plays/localhost.yml -t postgres

        docker-postgres-create -c postgres01 -d test
        docker-postgres-list -c postgres01 | grep "test"
        docker-postgres-shell -c postgres01 -d test -s 'CREATE TABLE example(id int);'
        docker-postgres-duplicate -c postgres01 -s test -t example
        docker-postgres-drop -c postgres01 -d test
        docker-postgres-size -c postgres01 -d example | grep "0 bytes"
        docker-postgres-backup -c postgres01 -d example
        docker-postgres-restore -c postgres01 -f /var/tmp/postgres01/example.sql -d test -i 
        docker-postgres-size -c postgres01 -d test | grep "0 bytes"
        mkdir -p tmp
        echo 'DROP TABLE example;' > tmp/drop.sql
        docker-postgres-patch -c postgres01 -d test -f ./tmp/drop.sql
        docker-postgres-rename -c postgres01 -s test -t example -r
    
        echo -e "\033[38;5;214mTEST\033[0m: Test odoo scripts"
        sudo ./roles/odoo_scripts/files/install

        # Start odoo container
        ansible-playbook -i test plays/localhost.yml -t odoo --skip-tags depends

        docker-odoo-init -c odoo01-prod -d odoo-prod -i contacts -w
        docker-odoo-rename -c odoo01-prod -s odoo-prod -t odoo
        docker-odoo-list -c odoo01-prod | grep "odoo$"
        docker-odoo-user -c odoo01-prod -d odoo -a list | grep "admin"
        docker-odoo-drop -c odoo01-int -d odoo-int
        docker-odoo-duplicate -c odoo01-prod -s odoo -t odoo-int
        docker-odoo-neutralize -c odoo01-int -d odoo-int
        docker exec odoo01-int git-clone-addons
        docker restart odoo01-int
        docker-odoo-init -c odoo01-int -d odoo-int -i partner_firstname
        docker-odoo-shell -c odoo01-int -d odoo-int -f -p 'env["res.partner"].browse(1).write({"name": "admint"})'
        [[ "admint" == $(docker-odoo-shell -c odoo01-int -d odoo-int -p 'print(env["res.partner"].browse(1).name)') ]]

        echo -e "\033[38;5;214mTEST\033[0m: Test docker volume scripts"
        docker-volume-backup -c postgres01 -v postgres01
        docker rm -f postgres01
        docker volume rm postgres01
        docker-volume-restore -v postgres01 -f /var/tmp/postgres01/postgres01.tar
        ansible-playbook -i test plays/localhost.yml -t postgres --skip-tags depends
        docker-postgres-list -c postgres01 | grep "test"
    fi
}

# LLM Commands

function update-with-llm() {
    if test -z "$1"; then 
        echo "\$1 is empty."; 
        exit 1; 
    fi

    # Get files from path
    FILES=$(find "$1" -type f \( -name "*.md" -o -name "*.yml" -o -name "*.yaml" -o -name "*.conf" -o -name "*.nginx" \) \
        | grep -vFf <(git -C "$1" ls-files --ignored --exclude-standard --others))
    echo -e "Loaded these files into prompt:\n\n$FILES\n"

    # Prompt task description
    if [ -z "$2" ]; then
        read -p "Enter the task description: " TASK_DESCRIPTION
    else
        TASK_DESCRIPTION="$2"
    fi

    # Prepare the files content for prompt
    FILE_CONTENTS=""
    for FILE in $FILES; do
        FILE_CONTENTS+="<<<$FILE>>>
$(cat "$FILE")

"
    done

    # Define prompt content
    mkdir -p tmp
    PROMPT_FILE="tmp/llm_update"
    echo -e "\nWrite prompt to $PROMPT_FILE."
    cat << EOF > "$PROMPT_FILE"
Look at the code files below and do the following:

$TASK_DESCRIPTION

Output all files that you need to change in full again, including your changes. 
In the same format as I provide the files below. Under no circumstances output any other text, 
no additional infos, no code formatting chars. Only the code in the given format.

Here are the files:

$FILE_CONTENTS
EOF

    # Run the llm command
    echo -e "Send prompt and wait for the response of the $LLM_MODEL LLM."
    RESULT=$(cat "$PROMPT_FILE" | llm -m "$LLM_MODEL")

    # Check if result is empty
    if [ -z "$RESULT" ]; then
        echo "No response from the model. Exiting."
        exit 1
    fi

    # Save the result to a file
    RESULT_FILE="tmp/llm_update_result"
    echo "$RESULT" > "$RESULT_FILE"
    echo -e "Saved response to $RESULT_FILE.\n"

    # Show a preview of result file
    less "$RESULT_FILE"

    # Ask for confirmation before updating files
    read -p "Do you want to apply these updates to the files? (y/n): " CONFIRM
    if [[ ! "$CONFIRM" =~ ^[Yy]$ ]]; then
        exit 0
    fi

    # Parse the response from the file and update the files directly
    echo -e "Parsing the response and updating files...\n"
    CURRENT_FILE=""
    while IFS= read -r LINE; do
        if [[ $LINE =~ ^"<<<"(.*)">>>" ]]; then
            CURRENT_FILE="${BASH_REMATCH[1]}"
            echo "Update file $CURRENT_FILE."
            > "$CURRENT_FILE"
        elif [[ -n $CURRENT_FILE ]]; then
            echo "$LINE" >> "$CURRENT_FILE"
        fi
    done < "$RESULT_FILE"
}

if declare -f "$1" > /dev/null; then
    "$1" "${@:2}"
else
    case "$1" in
        help)
            help task
            exit 1
            ;;
        install-node)
            npm install
            ;;
        dev-node)
            npm run dev
            ;;
        build-node)
            npm run build
            ;;
        serve-node-build)
            cd .vuepress/dist
            npx serve
            ;;
        lint)
            activate-venv
            ansible-lint --write --offline roles/*
            ;;
        source)
            activate-venv
            ;;
        *)
            echo "Unknown command: $1"
            help task
            exit 1
            ;;
    esac
fi