#!/bin/bash
# Copyright 2025-2026 Arm Limited and/or its affiliates.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

# Calling this script with one argument is equal to launching it in
# non-interactive mode. "$#" gives the number of positional arguments.
[ "$#" -eq 1 ] && is_script_interactive=1 || is_script_interactive=0

if [ $is_script_interactive -eq 1 ]; then
    RESET='\e[0m'
    RED='\e[31m'
    GREEN='\e[32m'
    YELLOW='\e[33m'
    BLUE='\e[34m'
fi

INFO="${BLUE}[INFO]${RESET}"
WARNING="${YELLOW}[WARNING]${RESET}"
ERROR="${RED}[ERROR]${RESET}"
SUCCESS="${GREEN}[SUCCESS]${RESET}"

DOCGEN_OUTPUTS=(
    "docs/source/backends/arm-ethos-u/arm-ethos-u-overview.md"
    "docs/source/backends/arm-ethos-u/arm-ethos-u-partitioner.md"
    "docs/source/backends/arm-ethos-u/arm-ethos-u-quantization.md"
    "docs/source/backends/arm-ethos-u/tutorials/ethos-u-getting-started.md"
    "docs/source/backends/arm-vgf/arm-vgf-overview.md"
    "docs/source/backends/arm-vgf/arm-vgf-partitioner.md"
    "docs/source/backends/arm-vgf/arm-vgf-quantization.md"
    "docs/source/backends/arm-vgf/tutorials/vgf-getting-started.md"
)

OP_SUPPORT_SCRIPT="backends/arm/scripts/docgen/generate_op_support.py"
VGF_OP_SUPPORT_OUTPUT="docs/source/backends/arm-vgf/VGF_op_support.md"
U55_OP_SUPPORT_OUTPUT="docs/source/backends/arm-ethos-u/U55_op_support.md"
U85_OP_SUPPORT_OUTPUT="docs/source/backends/arm-ethos-u/U85_op_support.md"

RUN_DOCGEN=0

is_docgen_trigger_file() {
    local file_path="$1"
    local output_path
    for output_path in "${DOCGEN_OUTPUTS[@]}"; do
        if [[ "$file_path" == "$output_path" ]]; then
            return 0
        fi
    done

    case "${file_path}" in
        backends/arm/common/arm_compile_spec.py|\
        backends/arm/ethosu/*.py|\
        backends/arm/quantizer/*.py|\
        backends/arm/scripts/docgen/*|\
        backends/arm/tosa/partitioner.py|\
        backends/arm/vgf/*.py|\
        examples/arm/ethos_u_minimal_example.ipynb|\
        examples/arm/vgf_minimal_example.ipynb)
            return 0
            ;;
    esac
    return 1
}

run_docgen_check() {
    local docgen_status=0

    echo -e "${INFO} Running Arm docgen"
    if ! python backends/arm/scripts/docgen/docgen.py; then
        echo -e "${ERROR} Failed to run Arm docgen"
        return 1
    fi

    local untracked_outputs
    untracked_outputs=$(git ls-files --others -- "${DOCGEN_OUTPUTS[@]}")
    if ! git diff --quiet HEAD -- "${DOCGEN_OUTPUTS[@]}" || \
        [[ -n "$untracked_outputs" ]]; then
        echo -e "${ERROR} Arm docgen updated generated documentation." >&2
        echo -e "${INFO} Review and commit the regenerated docs before pushing."
        git diff HEAD -- "${DOCGEN_OUTPUTS[@]}"
        if [[ -n "$untracked_outputs" ]]; then
            echo -e "${INFO} Untracked generated documentation:"
            echo "$untracked_outputs"
        fi
        docgen_status=1
    else
        echo -e "${SUCCESS} Arm docgen OK"
    fi
    return "$docgen_status"
}

run_op_support_checks() {
    local backend="$1"
    local backend_name="$2"
    local output="$3"
    local coverage_output
    local FAILED=0

    echo -e "${INFO} Generating ${backend_name} operator support documentation"

    if ! python "$OP_SUPPORT_SCRIPT" --backend "$backend"; then
        echo -e "${ERROR} Failed to generate ${backend_name} operator support documentation"
        FAILED=1
    elif ! git diff --quiet HEAD -- "$output"; then
        echo -e "${ERROR} ${backend_name} operator support documentation is out of date." >&2
        echo -e "${INFO} Review and commit the regenerated operator support list before pushing."
        git diff HEAD -- "$output"
        FAILED=1
    else
        echo -e "${SUCCESS} ${backend_name} operator support documentation OK"
    fi

    echo -e "${INFO} Checking ${backend_name} operator support coverage"

    if ! coverage_output=$(python "$OP_SUPPORT_SCRIPT" --backend "$backend" --check 2>&1); then
        echo "$coverage_output" >&2
        echo -e "${ERROR} ${backend_name} operator support coverage check failed" >&2
        FAILED=1
    else
        echo -e "${SUCCESS} ${backend_name} operator support coverage OK"
    fi
    return "$FAILED"
}

wait_public_api_validator() {
    if ! wait "${manifest_pid}"; then
        echo -e "${ERROR} Arm public API manifest validation failed"
        FAILED=1
    fi
}

# This list of imperative verbs was compiled from the entire list of Executorch
# commits. It should be fairly exhaustive, but add more verbs if you find one
# that's missing.
VERBS="Add|Fix|Update|Refactor|Improve|Remove|Change|Implement|Create|Modify|"\
"Enable|Integrate|Make|Support|Deprecate|Extend|Enhance|Convert|Rewrite|Unify|"\
"Optimize|Expand|Reorganize|Adjust|Streamline|Clarify|Introduce|Document|"\
"Polish|Standardize|Revise|Simplify|Restore|Resolve|Replace|Suppress|Migrate|"\
"Generate|Delete|Exclude|Register|Include|Upgrade|Validate|Verify|Refine|"\
"Reimplement|Patch|Sync|Revert|Fixup|Enhance|Append|Annotate|Disable|Emit|"\
"Handle|Ignore|Interpret|Instantiate|Invoke|Limit|Load|Modify|Permit|Print|"\
"Profile|Recalculate|Reconstruct|Redefine|Redesign|Reevaluate|Relocate|Remap|"\
"Render|Reposition|Request|Revert|Sanitize|Specify|Strengthen|Stub|Substitute|"\
"Tag|Tweak|Unify|Unlock|Unset|Use|Validate|Verify|Rename|Relax|Format|Don't|"\
"Consolidate|Review|Run|Parallelize"

# Remote branch
REMOTE=$(git rev-parse --abbrev-ref --symbolic-full-name @{u} 2>/dev/null)
if [ $is_script_interactive -eq 0 ]; then
    # Just use the one commit
    COMMITS=$(git rev-list HEAD -n 1)
elif [ -z "$REMOTE" ]; then
    echo -e "${WARNING} Could not find upstream branch to compare to."
    echo "Please specify the number of commits you are pushing."
    echo -n "Enter number of commits to check (default 1): " > /dev/tty
    read NUM_COMMITS < /dev/tty
    NUM_COMMITS=${NUM_COMMITS:-1} # Default to 1 if empty
    RANGE=$(git rev-list HEAD -n "$NUM_COMMITS")
    COMMITS=${RANGE}
elif [ "$(git rev-parse --abbrev-ref HEAD)" == "HEAD" ]; then
    echo -e "${WARNING} You're in a detached HEAD state."
    echo "Please specify the number of commits you are pushing."
    echo -n "Enter number of commits to check (default 1): " > /dev/tty
    read NUM_COMMITS < /dev/tty
    NUM_COMMITS=${NUM_COMMITS:-1} # Default to 1 if empty
    RANGE=$(git rev-list HEAD -n "$NUM_COMMITS")
    COMMITS=${RANGE}
else
    # Determine commits to check
    RANGE="$REMOTE..HEAD"
    COMMITS=$(git rev-list "$RANGE")
fi

if [ -z "$COMMITS" ]; then
    echo -e "${INFO} No new commits to check."
    exit 0
fi

check_output_dir=$(mktemp -d) || exit 1
trap 'rm -rf "${check_output_dir}"' EXIT
backends/arm/scripts/public_api_manifest/validate_all_public_api_manifests.sh \
    > "${check_output_dir}/manifest" 2>&1 &
manifest_pid=$!

for COMMIT in ${COMMITS}; do
    if [ -n "$REMOTE" ] && git merge-base --is-ancestor "$COMMIT" "$REMOTE"; then
        echo -e "${INFO} Skipping commit ${COMMIT} (already on $REMOTE)"
        continue
    fi

    # If commit header contains WIP, everything is ok
    git rev-list --format=%s --max-count=1 ${COMMIT} | grep -q WIP && \
         continue

    echo -e "${INFO} Checking commit ${COMMIT}"

    # lintrunner on latest patches.
    echo -e "${INFO} Lintrunner"
    MYPYPATH=./src/ lintrunner --revision ${COMMIT}^1
    if [[ $? != 0 ]]; then
        echo -e "${ERROR} Failed linting"
        FAILED=1
    else
        echo -e "${SUCCESS} Lintrunner OK"
    fi

    commit_files=$(git diff-tree --no-commit-id --name-only \
        --diff-filter=ACDMR ${COMMIT} -r)
    if [[ $RUN_DOCGEN -eq 0 ]]; then
        for committed_file in $commit_files; do
            if is_docgen_trigger_file "${committed_file}"; then
                RUN_DOCGEN=1
                break
            fi
        done
    fi
    # Check license headers
    # We do a simple check of if all committed headers contain
    # "$current_year Arm". This does not guarantee OK in ci but should be ok
    # most of the time.
    echo -e "${INFO} License check"

    # Check only files with content changes.
    # Skip pure moves/renames (R100) where file content is unchanged.
    license_files=()
    while IFS=$'\t' read -r status path1 path2; do
        case "$status" in
            A|M)
                license_files+=("$path1")
                ;;
            R100)
                ;;
            R*)
                license_files+=("$path2")
                ;;
        esac
    done < <(git diff-tree --no-commit-id --name-status --diff-filter=AMR -r -M ${COMMIT})

    current_year=$(date +%Y)
    failed_license_check=false
    for committed_file in "${license_files[@]}"; do
        # Skip files with certain extensions
        case "$committed_file" in
            *.md|*.md.in|*.json|*.yml|*.yaml|*.cmake|*.patch|.gitignore|*.bzl|BUCK|*/BUCK|TARGETS|*/TARGETS|*/generated/*)
                echo -e "${INFO} Skipping license check for ${committed_file} (excluded extension)"
                continue
                ;;
        esac

        file_header=$(head "$committed_file")
        arm_copyright_regex="Copyright .*Arm Limited and/or its affiliates"
        if ! echo "$file_header" | grep -Eqi "$arm_copyright_regex"; then
            echo -e "${WARNING} No Arm copyright header in ${committed_file}"\
            " (skipping license year check)"
            continue
        fi

        if ! echo "$file_header" | grep -q "$current_year Arm"; then
            echo -e "${ERROR} Header in $committed_file did not contain"\
            "'$current_year Arm'"
            failed_license_check=true
        else
            echo -e "${SUCCESS} $committed_file passed license check"
        fi
    done

    if [[ ${#license_files[@]} -eq 0 ]]; then
        echo -e "${INFO} No files with content changes to check"
    elif [[ $failed_license_check == true ]]; then
        FAILED=1
    else
        echo -e "${SUCCESS} All files passed license check"
    fi

    # Check commit message
    echo -e "${INFO} Checking commit message"
    COMMIT_MSG=$(git log -1 --format=%B "$COMMIT")

    SUBJECT=$(echo "$COMMIT_MSG" | head -n1)
    BODY=$(echo "$COMMIT_MSG" | tail -n +2)

    # Check subject length (72 chars), except for revert commits.
    is_revert_commit=false
    if [[ "$SUBJECT" =~ ^(Arm\ backend:\ )?Revert(\ |$) ]]; then
        is_revert_commit=true
    fi

    SUBJECT_MAX_LEN=72
    if [ ${#SUBJECT} -gt ${SUBJECT_MAX_LEN} ] && \
        [[ $is_revert_commit == false ]]; then
        echo -e "${ERROR} Subject exceeds ${SUBJECT_MAX_LEN} characters:"\
            "'${SUBJECT}'" >&2

        FAILED=1
    else
        echo -e "${SUCCESS} Commit message subject OK"
    fi

    # Check body line length (72 chars)
    BODY_MAX_LEN=72
    line_number=2 # Subject + 1 empty line
    failed_body_check=false
    while IFS= read -r line; do
        # Ignore URLs when checking line length
        line_without_urls=$(echo "$line" | sed -E 's_https?://[^[:space:]]+__g')
        if [ ${#line_without_urls} -gt ${BODY_MAX_LEN} ]; then
            echo -e "${ERROR} Line ${line_number} in body exceeds"\
                "${BODY_MAX_LEN} characters: '$line'" >&2

            failed_body_check=true
        fi

        ((line_number++))
    done <<< "$BODY"

    if [[ $failed_body_check == true ]]; then
        FAILED=1
    else
        echo -e "${SUCCESS} Commit message body OK"
    fi

    # Check for Signed-off-by
    if ! echo "$COMMIT_MSG" | grep -qE "^Signed-off-by: "; then
        echo -e "${ERROR} Commit message must contain a 'Signed-off-by'"\
            "footer." >&2

        FAILED=1
    fi

    # Determine whether all modifications are under backends/arm or examples/arm
    only_arm_changes=true
    for committed_file in $commit_files; do
        if [[ $committed_file != backends/arm/* && $committed_file != examples/arm/* ]]; then
            only_arm_changes=false
            break
        fi
    done

    # Determine if subject has 'Arm backend: ' prefix
    subject_prefix="Arm backend: "
    has_arm_backend_prefix=false
    subject_without_prefix="$SUBJECT"
    if [[ "$SUBJECT" == ${subject_prefix}* ]]; then
        has_arm_backend_prefix=true
        subject_without_prefix=${SUBJECT#"${subject_prefix}"}
    fi

    # Require 'Arm backend: ' prefix when the commit only touches Arm backend/examples files
    if [[ $only_arm_changes == true && $has_arm_backend_prefix == false && \
        $is_revert_commit == false ]]; then
        echo -e "${ERROR} Subject must start with 'Arm backend: ' when only backends/arm or examples/arm files are modified." >&2
        FAILED=1
    fi

    # Imperative mood check (after optional prefix removal)
    if [[ ! "$subject_without_prefix" =~ ^(${VERBS})(\ |$) ]]; then
        echo -e "${WARNING} Commit message should begin with an imperative verb." >&2
        if [ $is_script_interactive -eq 1 ]; then
            echo -n "There are warnings in your commit message. Do you want to"\
                "ignore the warning (y/N): " > /dev/tty

            read USER_INPUT < /dev/tty

            if [[ ! "$USER_INPUT" =~ ^[Yy]$ ]]; then
                FAILED=1
            fi
        fi
    fi

    echo "" # Newline to visually separate commit processing
done

docgen_pid=""
if [[ $RUN_DOCGEN -eq 1 ]]; then
    run_docgen_check > "${check_output_dir}/docgen" 2>&1 &
    docgen_pid=$!
else
    echo -e "${INFO} Skipping Arm docgen (no public API inputs changed)"
fi

run_op_support_checks "u55" "Ethos-U55" "$U55_OP_SUPPORT_OUTPUT" \
    > "${check_output_dir}/u55" 2>&1 &
u55_pid=$!
run_op_support_checks "u85" "Ethos-U85" "$U85_OP_SUPPORT_OUTPUT" \
    > "${check_output_dir}/u85" 2>&1 &
u85_pid=$!
run_op_support_checks "vgf" "VGF" "$VGF_OP_SUPPORT_OUTPUT" \
    > "${check_output_dir}/vgf" 2>&1 &
vgf_pid=$!

if [[ -n "$docgen_pid" ]]; then
    wait "$docgen_pid" || FAILED=1
    cat "${check_output_dir}/docgen"
fi
wait "$u55_pid" || FAILED=1
cat "${check_output_dir}/u55"
wait "$u85_pid" || FAILED=1
cat "${check_output_dir}/u85"
wait "$vgf_pid" || FAILED=1
cat "${check_output_dir}/vgf"
wait_public_api_validator
cat "${check_output_dir}/manifest"

if [[ $FAILED ]]; then
    echo -e "${INFO} Fix your commit message errors with"\
        "'git commit --amend' or 'git commit --fixup=<SHA>'"

    exit 1
else
    echo -e "${SUCCESS} All checks passed"
fi

exit 0
