#!/usr/bin/env bash
set -euo pipefail

# Script Name: ports
# Description: Enhanced port viewer with colors, filtering, and process info
# Usage: ports           # Show all ports (colorized)
#        ports -l        # Listening only (most common)
#        ports -p        # Show process names
#        ports 80        # Find what's on port 80
#        ports tcp       # TCP only
#        ports udp       # UDP only

VERSION="1.0.0"

# Colors
readonly GREEN='\033[0;32m'
readonly YELLOW='\033[1;33m'
readonly BLUE='\033[0;34m'
readonly RED='\033[0;31m'
readonly CYAN='\033[0;36m'
readonly BOLD='\033[1m'
readonly NC='\033[0m'

show_help() {
  echo -e "\033[1mports\033[0m - Enhanced Port Viewer v${VERSION}"
  echo
  echo -e "\033[1mUSAGE:\033[0m"
  echo "  ports [OPTIONS] [PORT|PROTOCOL]"
  echo
  echo -e "\033[1mOPTIONS:\033[0m"
  echo -e "  \033[0;36m-l, --listen\033[0m      Show listening ports only (default)"
  echo -e "  \033[0;36m-a, --all\033[0m         Show all connections"
  echo -e "  \033[0;36m-p, --process\033[0m     Show process names/PIDs"
  echo -e "  \033[0;36m-n, --numeric\033[0m     Don't resolve hostnames"
  echo -e "  \033[0;36m-h, --help\033[0m        Show this help message"
  echo
  echo -e "\033[1mFILTERS:\033[0m"
  echo "  ports 80          # Show what's on port 80"
  echo "  ports tcp         # TCP connections only"
  echo "  ports udp         # UDP connections only"
  echo "  ports 8000-9000   # Port range"
  echo
  echo -e "\033[1mEXAMPLES:\033[0m"
  echo "  ports             # Listening ports (colorized)"
  echo "  ports -p          # With process info"
  echo "  ports -a          # All connections"
  echo "  ports 443         # What's on HTTPS port"
  echo "  ports tcp -p      # TCP with processes"
  echo
  echo -e "\033[1mCOLOR LEGEND:\033[0m"
  echo -e "  \033[0;32mLISTEN\033[0m      - Listening for connections"
  echo -e "  \033[0;34mESTABLISHED\033[0m - Active connection"
  echo -e "  \033[1;33mTIME_WAIT\033[0m   - Connection closing"
  echo -e "  \033[0;31mCLOSE_WAIT\033[0m  - Waiting to close"
}

# Colorize state
colorize_state() {
  local state=$1
  case "$state" in
    LISTEN)
      echo -e "${GREEN}${state}${NC}"
      ;;
    ESTABLISHED)
      echo -e "${BLUE}${state}${NC}"
      ;;
    TIME_WAIT|CLOSE_WAIT)
      echo -e "${YELLOW}${state}${NC}"
      ;;
    SYN_SENT|SYN_RECV)
      echo -e "${CYAN}${state}${NC}"
      ;;
    *)
      echo "$state"
      ;;
  esac
}

# Check if port is "interesting" (not in common_ports)
is_unusual_port() {
  local port=$1
  local common_ports=(20 21 22 23 25 53 80 110 143 443 465 587 993 995 3306 5432 6379 8080 8443)

  for p in "${common_ports[@]}"; do
    if [[ "$port" == "$p" ]]; then
      return 1
    fi
  done
  return 0
}

# Parse arguments
show_listen_only=true
show_process=false
numeric=true
filter_proto=""
filter_port=""

while [[ $# -gt 0 ]]; do
  case $1 in
    -l|--listen)
      show_listen_only=true
      shift
      ;;
    -a|--all)
      show_listen_only=false
      shift
      ;;
    -p|--process)
      show_process=true
      shift
      ;;
    -n|--numeric)
      numeric=true
      shift
      ;;
    -h|--help)
      show_help
      exit 0
      ;;
    tcp|TCP)
      filter_proto="tcp"
      shift
      ;;
    udp|UDP)
      filter_proto="udp"
      shift
      ;;
    [0-9]*)
      filter_port="$1"
      shift
      ;;
    *)
      echo -e "${RED}Error:${NC} Unknown option: $1" >&2
      echo "Run 'ports --help' for usage information" >&2
      exit 1
      ;;
  esac
done

# Build ss command
ss_cmd="ss -tuln"

if [[ "$show_listen_only" == "false" ]]; then
  ss_cmd="ss -tun"
fi

if [[ "$show_process" == "true" ]]; then
  ss_cmd="sudo ss -tulnp"
  if [[ "$show_listen_only" == "false" ]]; then
    ss_cmd="sudo ss -tunp"
  fi
fi

# Execute and format
output=$($ss_cmd)

# Header
echo -e "${BOLD}${CYAN}Active Ports${NC}"
echo -e "${BOLD}────────────────────────────────────────────────────────────${NC}"

# Parse and colorize output
echo "$output" | awk -v show_proc="$show_process" -v filter_proto="$filter_proto" -v filter_port="$filter_port" '
NR==1 { next }  # Skip header from ss

{
  proto = $1
  state = $2
  local = $5
  peer = $6
  process = ""

  # Extract process info if available (last field)
  if (show_proc == "true" && NF >= 7) {
    for (i=7; i<=NF; i++) {
      process = process $i " "
    }
  }

  # Filter by protocol
  if (filter_proto != "" && tolower(proto) !~ tolower(filter_proto)) next

  # Extract port from local address
  split(local, parts, ":")
  port = parts[length(parts)]

  # Filter by port
  if (filter_port != "" && port != filter_port) next

  # Print formatted line
  printf "%-6s %-12s %-25s %-25s", proto, state, local, peer

  if (process != "") {
    printf " %s", process
  }

  printf "\n"
}
' | while IFS= read -r line; do
  # Colorize based on state
  if [[ "$line" =~ LISTEN ]]; then
    echo -e "$line" | sed "s/LISTEN/${GREEN}LISTEN${NC}/"
  elif [[ "$line" =~ ESTABLISHED ]]; then
    echo -e "$line" | sed "s/ESTABLISHED/${BLUE}ESTABLISHED${NC}/"
  elif [[ "$line" =~ TIME_WAIT ]]; then
    echo -e "$line" | sed "s/TIME_WAIT/${YELLOW}TIME_WAIT${NC}/"
  elif [[ "$line" =~ CLOSE_WAIT ]]; then
    echo -e "$line" | sed "s/CLOSE_WAIT/${RED}CLOSE_WAIT${NC}/"
  else
    echo "$line"
  fi
done

# Summary
echo
echo -e "${BOLD}${CYAN}Summary:${NC}"
total=$(echo "$output" | wc -l)
echo "  Total connections: $((total - 1))"

if [[ "$show_listen_only" == "true" ]]; then
  echo -e "  ${GREEN}Tip:${NC} Use 'ports -a' to see all connections"
fi

if [[ "$show_process" == "false" ]]; then
  echo -e "  ${GREEN}Tip:${NC} Use 'ports -p' to see process information"
fi