#!/bin/bash
# Copyright (c) Huawei Technologies Co., Ltd. 2026-2026. All rights reserved.
# ============================================================================
# KDC Workspace Unit Test Coverage Report
# ============================================================================
# Collects and displays line / function / branch coverage for all workspace
# crates using cargo-llvm-cov (nightly toolchain required for branch coverage).
#
# Prerequisites:
#   1. rustup install nightly && rustup component add llvm-tools-preview --toolchain nightly
#   2. cargo install cargo-llvm-cov
#   3. jq (JSON processor)
#
# Usage:
#   bash scripts/coverage.sh            # Run coverage and show console report
#   bash scripts/coverage.sh --html     # Also generate HTML report
#   bash scripts/coverage.sh --json     # Also keep raw JSON at target/coverage/coverage.json
# ============================================================================

set -euo pipefail

# ── Colors ────────────────────────────────────────────────────────────────────
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
BOLD='\033[1m'
DIM='\033[2m'
RESET='\033[0m'

# ── Project root ─────────────────────────────────────────────────────────────
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
COVERAGE_DIR="${PROJECT_ROOT}/target/coverage"
COVERAGE_JSON="${COVERAGE_DIR}/coverage.json"

# ── Parse arguments ──────────────────────────────────────────────────────────
GENERATE_HTML=true
KEEP_JSON=true
for arg in "$@"; do
    case "${arg}" in
        --html) GENERATE_HTML=true ;;
        --json) KEEP_JSON=true ;;
        --help|-h)
            echo "Usage: bash scripts/coverage.sh [--html] [--json]"
            echo ""
            echo "  --html    Generate HTML coverage report"
            echo "  --json    Keep raw JSON at target/coverage/coverage.json"
            echo "  --help    Show this help message"
            exit 0
            ;;
        *)
            echo "Unknown option: ${arg}"
            echo "Usage: bash scripts/coverage.sh [--html] [--json]"
            exit 1
            ;;
    esac
done

# ── Helper functions ─────────────────────────────────────────────────────────
info()  { echo -e "${CYAN}[INFO]${RESET} $*"; }
warn()  { echo -e "${YELLOW}[WARN]${RESET} $*"; }
error() { echo -e "${RED}[ERROR]${RESET} $*" >&2; }

# ── Prerequisite checks ──────────────────────────────────────────────────────
check_prerequisites() {
    local missing=()

    if ! command -v cargo-llvm-cov &>/dev/null; then
        missing+=("cargo-llvm-cov  → cargo install cargo-llvm-cov")
    fi

    if ! command -v jq &>/dev/null; then
        missing+=("jq  → yum install jq / apt-get install jq")
    fi

    if ! rustup toolchain list | grep -q "nightly"; then
        missing+=("nightly toolchain  → rustup install nightly")
    fi

    if ! rustup component list --toolchain nightly | grep -q "llvm-tools.*installed"; then
        missing+=("llvm-tools-preview (nightly)  → rustup component add llvm-tools-preview --toolchain nightly")
    fi

    if [ ${#missing[@]} -ne 0 ]; then
        error "Missing prerequisites:"
        for item in "${missing[@]}"; do
            echo "  - ${item}"
        done
        exit 1
    fi
}

# ── Run coverage collection ──────────────────────────────────────────────────
collect_coverage() {
    mkdir -p "${COVERAGE_DIR}"

    info "Running coverage (nightly, branch-aware, single-threaded, incremental, lib-only)..."
    # Use --lib to exclude binary targets (main.rs is untestable).
    # Binary compilation links lib functions but never calls them, which
    # would dilute coverage metrics with false "uncovered" entries.
    #
    # Step 1: Run tests with --no-report to collect profiling data only (no report generation).
    # Step 2: Use `report` subcommand to generate JSON/HTML from the same data without re-running tests.
    if cargo +nightly llvm-cov --lib --workspace --branch \
        --no-report \
        -- --test-threads=1 2>&1; then
        :
    else
        local rc=$?
        error "Coverage collection failed (exit code ${rc})"
        exit ${rc}
    fi

    # Generate JSON report from collected profiling data
    cargo +nightly llvm-cov report --branch \
        --json --output-path "${COVERAGE_JSON}" 2>/dev/null || true

    # Generate HTML report from collected profiling data
    if [ "${GENERATE_HTML}" = true ]; then
        info "Generating HTML report..."
        cargo +nightly llvm-cov report --branch --html \
            --output-dir "${COVERAGE_DIR}" 2>/dev/null || true

        # Append custom CSS to make coverage colors more intuitive:
        # green background = covered, red background = uncovered, gray = skipped.
        local css_file="${COVERAGE_DIR}/html/style.css"
        if [ -f "${css_file}" ]; then
            cat >> "${css_file}" <<'STYLE_EOF'

/* ── Custom: intuitive background colors for coverage ── */
td.covered-line { background-color: #e6f4ea; color: #1e7e34; font-weight: bold; }
td.uncovered-line { background-color: #fce8e6; color: #c62828; font-weight: bold; }
td.skipped-line { background-color: #f5f5f5; color: #999; }
STYLE_EOF
        fi

        info "HTML report: ${COVERAGE_DIR}/html/index.html"
    fi

    if [ "${KEEP_JSON}" = false ]; then
        # JSON is still needed for parsing below; clean up at end
        trap "rm -f '${COVERAGE_JSON}'" EXIT
    fi
}

# ── Format percentage with color ─────────────────────────────────────────────
# Usage: format_pct <percentage>
format_pct() {
    # Normalize to exactly 2 decimal places for consistent column alignment
    local pct
    pct=$(printf "%.2f" "${1}" 2>/dev/null || echo "0.00")
    if (( $(echo "${pct} == 100" | bc -l) )); then
        pct="100.0"
    fi
    if (( $(echo "${pct} >= 80" | bc -l) )); then
        echo -e "${GREEN}${pct}%${RESET}"
    elif (( $(echo "${pct} >= 60" | bc -l) )); then
        echo -e "${YELLOW}${pct}%${RESET}"
    else
        echo -e "${RED}${pct}%${RESET}"
    fi
}

# ── Classify file path into crate name ───────────────────────────────────────
# Usage: classify_crate <filepath>
classify_crate() {
    local filepath="$1"
    case "${filepath}" in
        */kdc_agent/*)                  echo "kdc_agent" ;;
        */kdc_proxy/*)                  echo "kdc_proxy" ;;
        *)                              echo "other" ;;
    esac
}

# ── Build coverage report from JSON ──────────────────────────────────────────
build_report() {
    if [ ! -f "${COVERAGE_JSON}" ]; then
        error "Coverage JSON not found: ${COVERAGE_JSON}"
        exit 1
    fi

    # ── Table header ─────────────────────────────────────────────────────
    echo ""
    echo -e "${BOLD}╔══════════════════════════════════════════════════════════════════════════════════════╗${RESET}"
    echo -e "${BOLD}║                           KDC Unit Test Coverage Report                              ║${RESET}"
    echo -e "${BOLD}╠══════════════════════════════════════════════════════════════════════════════════════╣${RESET}"
    echo -e "${BOLD}║  Crate            │  Lines                │  Functions          │  Branches          ║${RESET}"
    echo -e "${BOLD}╠══════════════════════════════════════════════════════════════════════════════════════╣${RESET}"

    # ── Per-crate rows ───────────────────────────────────────────────────
    # Extract per-file data and aggregate by crate using jq
    local crate_data
    crate_data=$(jq -r '
        .data[0].files[] |
        .filename as $file |
        .summary |
        {
            crate: (if $file | test("kdc_agent") then "kdc_agent"
                    elif $file | test("kdc_proxy") then "kdc_proxy"
                    else "other" end),
            lines_count: .lines.count,
            lines_covered: .lines.covered,
            func_count: .functions.count,
            func_covered: .functions.covered,
            branch_count: .branches.count,
            branch_covered: .branches.covered
        }
    ' "${COVERAGE_JSON}")

    # Aggregate totals per crate
    local aggregated
    aggregated=$(echo "${crate_data}" | jq -s 'group_by(.crate) | map({
        crate: .[0].crate,
        lines_count:   (map(.lines_count)    | add // 0),
        lines_covered: (map(.lines_covered)  | add // 0),
        func_count:    (map(.func_count)     | add // 0),
        func_covered:  (map(.func_covered)   | add // 0),
        branch_count:  (map(.branch_count)   | add // 0),
        branch_covered:(map(.branch_covered) | add // 0)
    })')

    # Sort order: kdc_agent, kdc_proxy
    local order='["kdc_agent", "kdc_proxy"]'
    local sorted
    sorted=$(echo "${aggregated}" | jq --argjson order "${order}" '
        sort_by(.crate as $c | $order | index($c) // 999)
    ')

    # Accumulate workspace totals
    local ws_lines_count=0 ws_lines_covered=0
    local ws_func_count=0 ws_func_covered=0
    local ws_branch_count=0 ws_branch_covered=0

    # Print each crate row
    echo "${sorted}" | jq -c '.[]' | while read -r row; do
        local crate_name lc lcov fc fcov bc bcov
        crate_name=$(echo "${row}" | jq -r '.crate')
        lc=$(echo "${row}" | jq '.lines_count')
        lcov=$(echo "${row}" | jq '.lines_covered')
        fc=$(echo "${row}" | jq '.func_count')
        fcov=$(echo "${row}" | jq '.func_covered')
        bc=$(echo "${row}" | jq '.branch_count')
        bcov=$(echo "${row}" | jq '.branch_covered')

        local lpct fpct bpct
        lpct=$(echo "scale=2; ${lcov} * 100 / ${lc}" | bc -l 2>/dev/null || echo "0")
        fpct=$(echo "scale=2; ${fcov} * 100 / ${fc}" | bc -l 2>/dev/null || echo "0")
        bpct=$(echo "scale=2; ${bcov} * 100 / ${bc}" | bc -l 2>/dev/null || echo "0")

        # Format: right-align numbers, fixed-width columns
        printf "║  %-16s │  %5d/%-5d %8s   │  %4d/%-4d %8s   │  %4d/%-4d %8s  ║\n" \
            "${crate_name}" \
            "${lcov}" "${lc}" "$(format_pct "${lpct}")" \
            "${fcov}" "${fc}" "$(format_pct "${fpct}")" \
            "${bcov}" "${bc}" "$(format_pct "${bpct}")"
    done

    # ── Workspace totals ─────────────────────────────────────────────────
    local totals
    totals=$(jq -r '
        .data[0].totals |
        {
            lc: .lines.count,       lcov: .lines.covered,
            fc: .functions.count,   fcov: .functions.covered,
            bc: .branches.count,    bcov: .branches.covered
        }
    ' "${COVERAGE_JSON}")

    local tlc tlcov tfc tfcov tbc tbcov
    tlc=$(echo "${totals}" | jq '.lc')
    tlcov=$(echo "${totals}" | jq '.lcov')
    tfc=$(echo "${totals}" | jq '.fc')
    tfcov=$(echo "${totals}" | jq '.fcov')
    tbc=$(echo "${totals}" | jq '.bc')
    tbcov=$(echo "${totals}" | jq '.bcov')

    local tlpct tfpct tbpct
    tlpct=$(echo "scale=2; ${tlcov} * 100 / ${tlc}" | bc -l 2>/dev/null || echo "0")
    tfpct=$(echo "scale=2; ${tfcov} * 100 / ${tfc}" | bc -l 2>/dev/null || echo "0")
    tbpct=$(echo "scale=2; ${tbcov} * 100 / ${tbc}" | bc -l 2>/dev/null || echo "0")

    echo -e "${BOLD}╠══════════════════════════════════════════════════════════════════════════════════════╣${RESET}"
    printf "║  ${BOLD}%-16s${RESET} │  ${BOLD}%5d/%-5d %7s${RESET}   │  ${BOLD}%4d/%-4d %7s${RESET}   │  ${BOLD}%4d/%-4d %7s${RESET}  ║\n" \
        "TOTAL" \
        "${tlcov}" "${tlc}" "$(format_pct "${tlpct}")" \
        "${tfcov}" "${tfc}" "$(format_pct "${tfpct}")" \
        "${tbcov}" "${tbc}" "$(format_pct "${tbpct}")"
    echo -e "${BOLD}╚══════════════════════════════════════════════════════════════════════════════════════╝${RESET}"

    echo ""

    # ── Per-file detail ─────────────────────────────────────────────────
    echo -e "${BOLD}Per-File Coverage Detail${RESET}"
    echo -e "${BOLD}─────────────────────────────────────────────────────────────────────────────${RESET}"

    jq -r '
        def fpct:
            (. * 100 | . * 100 + 0.5 | floor | . / 100) as $v |
            ($v | tostring | split(".") | .[0] + "." + (.[1] + "00")[0:2]) as $r |
            if $r == "100.00" then "100.0%" else $r + "%" end;
        .data[0].files[] |
        .filename as $file |
        .summary |
        [
            (.lines.covered | tostring),
            (.lines.count   | tostring),
            (if .lines.count > 0 then ((.lines.covered / .lines.count) | fpct) else " 0.00%" end),
            (if .functions.count > 0 then ((.functions.covered / .functions.count) | fpct) else "   N/A" end),
            (if .branches.count > 0 then ((.branches.covered / .branches.count) | fpct) else "   N/A" end),
            ($file | split("/") | .[-3:] | join("/"))
        ] | @tsv
    ' "${COVERAGE_JSON}" | while IFS=$'\t' read -r lcov lc lpct fpct bpct fname; do
        local color="${RESET}"
        local lpct_raw="${lpct%\%}"  # strip trailing % for numeric comparison
        if [ "${lc}" -gt 0 ] 2>/dev/null; then
            local pct_val
            pct_val=$(echo "scale=0; ${lpct_raw} >= 80" | bc 2>/dev/null || echo "0")
            if [ "${pct_val}" -eq 1 ] 2>/dev/null; then
                color="${GREEN}"
            else
                pct_val=$(echo "scale=0; ${lpct_raw} >= 60" | bc 2>/dev/null || echo "0")
                if [ "${pct_val}" -eq 1 ] 2>/dev/null; then
                    color="${YELLOW}"
                else
                    color="${RED}"
                fi
            fi
        fi
        local formatted
        formatted=$(printf "  %5s/%-5s %6s  %6s  %6s  %s" "${lcov}" "${lc}" "${lpct}" "${fpct}" "${bpct}" "${fname}")
        echo -e "${color}${formatted}${RESET}"
    done

    echo ""
    echo -e "${DIM}Legend: GREEN >= 80%  YELLOW >= 60%  RED < 60%${RESET}"
    echo ""
}

# ── Main ─────────────────────────────────────────────────────────────────────
main() {
    check_prerequisites
    collect_coverage
    build_report
    info "Coverage report generated successfully."
}

main "$@"