#!/usr/bin/env bash

#########################################################################################
#### port-checker.sh — Check host:port pairs and report open/closed/timeout status   ####
#### Accepts targets from CLI args, a file, or stdin                                 ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.00                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./port-checker.sh host1:port1 host2:port2                                     ####
####   ./port-checker.sh --file targets.txt                                          ####
####   echo "host:port" | ./port-checker.sh                                          ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
TIMEOUT="${PORT_CHECK_TIMEOUT:-3}"
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"
JSON_OUTPUT="${JSON_OUTPUT:-false}"
INPUT_FILE=""

# ── State ─────────────────────────────────────────────────────────────
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME
TARGETS=()
COUNT_OPEN=0
COUNT_CLOSED=0
COUNT_TIMEOUT=0
COUNT_TOTAL=0

# ── Colors ────────────────────────────────────────────────────────────
setup_colors() {
    if [[ "$COLOR" == "never" ]]; then
        RED="" GREEN="" YELLOW="" BOLD="" DIM="" RESET=""
        return
    fi
    if [[ "$COLOR" == "always" ]] || [[ -t 1 ]]; then
        RED='\033[0;31m'
        GREEN='\033[0;32m'
        YELLOW='\033[0;33m'
        BOLD='\033[1m'
        DIM='\033[2m'
        RESET='\033[0m'
    else
        RED="" GREEN="" YELLOW="" BOLD="" DIM="" RESET=""
    fi
}

# ── Logging ───────────────────────────────────────────────────────────
log()     { echo -e "${DIM}[INFO]${RESET} $*"; }
warn()    { echo -e "${YELLOW}[WARN]${RESET} $*" >&2; }
err()     { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
verbose() { if [[ "$VERBOSE" == "true" ]]; then echo -e "${DIM}[DEBUG]${RESET} $*"; fi; }

# ── Helpers ───────────────────────────────────────────────────────────
check_port() {
    local host="$1"
    local port="$2"
    local status="closed"
    local latency="-"

    verbose "Checking ${host}:${port} (timeout ${TIMEOUT}s)"

    local start_ns end_ns elapsed_ms
    start_ns=$(date +%s%N)

    if command -v nc &>/dev/null; then
        if nc -z -w "$TIMEOUT" "$host" "$port" &>/dev/null; then
            status="open"
        fi
    elif command -v ncat &>/dev/null; then
        if ncat -z -w "$TIMEOUT" "$host" "$port" &>/dev/null; then
            status="open"
        fi
    else
        if (echo >/dev/tcp/"$host"/"$port") 2>/dev/null; then
            status="open"
        fi
    fi

    end_ns=$(date +%s%N)
    elapsed_ms=$(( (end_ns - start_ns) / 1000000 ))

    if [[ "$status" == "closed" && "$elapsed_ms" -ge $(( TIMEOUT * 1000 )) ]]; then
        status="timeout"
    fi

    if [[ "$status" != "closed" || "$elapsed_ms" -lt $(( TIMEOUT * 1000 )) ]]; then
        latency="${elapsed_ms}ms"
    fi

    COUNT_TOTAL=$((COUNT_TOTAL + 1))
    case "$status" in
        open)    COUNT_OPEN=$((COUNT_OPEN + 1)) ;;
        closed)  COUNT_CLOSED=$((COUNT_CLOSED + 1)) ;;
        timeout) COUNT_TIMEOUT=$((COUNT_TIMEOUT + 1)) ;;
    esac

    if [[ "$JSON_OUTPUT" == "true" ]]; then
        local json_comma=""
        if [[ "$COUNT_TOTAL" -gt 1 ]]; then
            json_comma=","
        fi
        printf '%s{"host":"%s","port":%s,"status":"%s","latency":"%s"}' \
            "$json_comma" "$host" "$port" "$status" "$latency"
    else
        local color
        case "$status" in
            open)    color="$GREEN" ;;
            closed)  color="$RED" ;;
            timeout) color="$YELLOW" ;;
            *)       color="" ;;
        esac
        printf "  %-30s %-8s %b%-10s%b %s\n" "$host" "$port" "$color" "$status" "$RESET" "$latency"
    fi
}

load_targets_from_file() {
    local file="$1"
    if [[ ! -f "$file" ]]; then
        err "File not found: $file"
        exit 1
    fi
    while IFS= read -r line; do
        line="${line%%#*}"
        line="${line// /}"
        [[ -z "$line" ]] && continue
        TARGETS+=("$line")
    done < "$file"
}

load_targets_from_stdin() {
    while IFS= read -r line; do
        line="${line%%#*}"
        line="${line// /}"
        [[ -z "$line" ]] && continue
        TARGETS+=("$line")
    done
}

# ══════════════════════════════════════════════════════════════════════
#                       USAGE
# ══════════════════════════════════════════════════════════════════════

usage() {
    cat <<EOF
${SCRIPT_NAME} — Check host:port pairs and report open/closed/timeout status

USAGE:
  ${SCRIPT_NAME} [OPTIONS] host1:port1 [host2:port2 ...]
  ${SCRIPT_NAME} --file targets.txt
  echo "host:port" | ${SCRIPT_NAME}

OPTIONS:
  --file FILE           Read targets from file (one host:port per line)
  --timeout SECONDS     Connection timeout (default: ${TIMEOUT})
  --json                Output results in JSON format
  --verbose             Enable debug output
  --no-color            Disable colored output
  --help                Show this help

ENVIRONMENT VARIABLES:
  PORT_CHECK_TIMEOUT    Connection timeout in seconds (default: 3)
  COLOR                 Color mode: auto, always, never (default: auto)

EXAMPLES:
  # Check specific hosts
  ./port-checker.sh google.com:443 github.com:22

  # Read from file
  ./port-checker.sh --file targets.txt

  # From stdin
  echo -e "google.com:443\ngithub.com:22" | ./port-checker.sh

  # JSON output with custom timeout
  ./port-checker.sh --json --timeout 5 myhost:8080

  # Pipe-friendly
  ./port-checker.sh --no-color myhost:80 myhost:443
EOF
}

# ══════════════════════════════════════════════════════════════════════
#                       ARGUMENT PARSING
# ══════════════════════════════════════════════════════════════════════

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --file)
                INPUT_FILE="$2"; shift 2 ;;
            --timeout)
                TIMEOUT="$2"; shift 2 ;;
            --json)
                JSON_OUTPUT="true"; shift ;;
            --verbose)
                VERBOSE="true"; shift ;;
            --no-color)
                COLOR="never"; shift ;;
            --help|-h)
                setup_colors
                usage
                exit 0 ;;
            -*)
                err "Unknown option: $1"
                echo "Run ${SCRIPT_NAME} --help for usage" >&2
                exit 1 ;;
            *)
                TARGETS+=("$1"); shift ;;
        esac
    done
}

# ══════════════════════════════════════════════════════════════════════
#                       MAIN
# ══════════════════════════════════════════════════════════════════════

main() {
    parse_args "$@"
    setup_colors

    # Load targets from file if specified
    if [[ -n "$INPUT_FILE" ]]; then
        load_targets_from_file "$INPUT_FILE"
    fi

    # Load from stdin if no targets yet and stdin is not a terminal
    if [[ ${#TARGETS[@]} -eq 0 ]] && ! [[ -t 0 ]]; then
        load_targets_from_stdin
    fi

    # Validate we have targets
    if [[ ${#TARGETS[@]} -eq 0 ]]; then
        err "No targets specified"
        echo "Run ${SCRIPT_NAME} --help for usage" >&2
        exit 1
    fi

    if [[ "$JSON_OUTPUT" == "true" ]]; then
        printf '{"results":['
    else
        echo ""
        echo -e "${BOLD}Port Checker${RESET}"
        echo -e "${DIM}Timeout: ${TIMEOUT}s${RESET}"
        echo ""
        printf "  ${BOLD}%-30s %-8s %-10s %s${RESET}\n" "HOST" "PORT" "STATUS" "LATENCY"
        printf "  %s\n" "$(printf '%.0s─' {1..62})"
    fi

    for target in "${TARGETS[@]}"; do
        local host port
        if [[ "$target" == *:* ]]; then
            host="${target%%:*}"
            port="${target##*:}"
        else
            warn "Invalid format (expected host:port): $target"
            continue
        fi

        if [[ -z "$host" || -z "$port" ]]; then
            warn "Invalid target: $target"
            continue
        fi

        check_port "$host" "$port"
    done

    if [[ "$JSON_OUTPUT" == "true" ]]; then
        printf '],"summary":{"total":%d,"open":%d,"closed":%d,"timeout":%d}}\n' \
            "$COUNT_TOTAL" "$COUNT_OPEN" "$COUNT_CLOSED" "$COUNT_TIMEOUT"
    else
        echo ""
        echo -e "  ${BOLD}Summary${RESET}"
        printf "  %-14s %d\n" "Total checked:" "$COUNT_TOTAL"
        printf "  %-14s %b%d%b\n" "Open:" "$GREEN" "$COUNT_OPEN" "$RESET"
        printf "  %-14s %b%d%b\n" "Closed:" "$RED" "$COUNT_CLOSED" "$RESET"
        printf "  %-14s %b%d%b\n" "Timeout:" "$YELLOW" "$COUNT_TIMEOUT" "$RESET"
        echo ""
    fi
}

main "$@"