#!/usr/bin/env bash

# codeview - A tool to visualize codebases for LLM interactions
# Version: 1.0.0

# Colors
USE_COLORS=true
RED_BOLD='\033[1;31m'
GREEN='\033[0;32m'
BLUE='\033[0;34m'
GRAY='\033[0;90m'
YELLOW='\033[0;33m'
RESET='\033[0m'

# Default values
OUTPUT_FORMAT="text"
INCLUDE_PATTERNS=("*.py" "*.md" "*.js" "*.html" "*.css" "*.json" "*.yaml" "*.yml")
EXCLUDE_DIRS=("myenv" "venv" ".venv" "node_modules" ".git" "__pycache__" ".pytest_cache" "build" "dist")
EXCLUDE_FILES=("*.pyc" "*.pyo" "*.pyd" "*.so" "*.dll" "*.class" "*.egg-info" "*.egg")
MAX_DEPTH="-1"
SHOW_TREE=true
SHOW_FILES=true
SHOW_LINE_NUMBERS=false
OUTPUT_FILE=""
SEARCH_PATTERN=""
INCLUDE_DIRS=()
TEMP_OUTPUT=""

print_usage() {
  local g=$GREEN r=$RESET y=$YELLOW
  $USE_COLORS || g="" r="" y=""
  
  printf "%bcodeview%b - A tool to visualize codebases for LLM interactions\n" "$g" "$r"
  printf "\nUsage: %bcodeview [options]%b\n\n" "$y" "$r"
  printf "Options:\n"
  printf "  -h, --help                 Show this help message\n"
  printf "  -i, --include PATTERN      File patterns to include (can be used multiple times)\n"
  printf "  -e, --exclude-dir DIR      Directories to exclude (can be used multiple times)\n"
  printf "  -x, --exclude-file PATTERN File patterns to exclude (can be used multiple times)\n"
  printf "  -d, --max-depth DEPTH      Maximum directory depth to traverse\n"
  printf "  -t, --no-tree              Don't show directory tree\n"
  printf "  -f, --no-files             Don't show file contents\n"
  printf "  -n, --line-numbers         Show line numbers in file contents\n"
  printf "  -o, --output FILE          Write output to file instead of stdout\n"
  printf "  -s, --search PATTERN       Only include files containing the pattern\n"
  printf "  -p, --path DIR             Include specific directory (can be used multiple times)\n"
  printf "  -m, --format FORMAT        Output format: text (default), markdown, json\n"
  printf "\nExamples:\n"
  printf "  codeview                                  # Show all code files in current directory\n"
  printf "  codeview -i \"*.py\" -i \"*.js\"           # Only show Python and JavaScript files\n"
  printf "  codeview -e node_modules -e .git         # Exclude node_modules and .git directories\n"
  printf "  codeview -d 2                            # Only traverse 2 directory levels deep\n"
  printf "  codeview -s \"def main\"                  # Only show files containing 'def main'\n"
  printf "  codeview -p src/models -p src/utils      # Only include specific directories\n"
  printf "  codeview -m markdown -o codebase.md      # Output in markdown format to a file\n"
}

command_exists() {
  command -v "$1" >/dev/null 2>&1
}

generate_tree() {
  local exclude_args=()
  local b=$BLUE r=$RESET y=$YELLOW
  $USE_COLORS || b="" r="" y=""
  
  if ! command_exists tree; then
    local rb=$RED_BOLD
    $USE_COLORS || rb=""
    printf "%bWarning: 'tree' command not found. Directory structure will not be shown.%b\n\n" "$rb" "$r"
    printf "%bInstall it with: apt-get install tree (Debian/Ubuntu) or brew install tree (macOS)%b\n\n" "$y" "$r"
    return
  fi
  for dir in "${EXCLUDE_DIRS[@]}"; do exclude_args+=("-I" "$dir"); done
  for file in "${EXCLUDE_FILES[@]}"; do exclude_args+=("-I" "$file"); done
  local depth_arg=()
  [ "$MAX_DEPTH" -ge 0 ] && depth_arg=("-L" "$MAX_DEPTH")
  printf "%bDirectory Structure:%b\n\n" "$b" "$r"
  if [ ${#INCLUDE_DIRS[@]} -eq 0 ]; then
    tree "${depth_arg[@]}" "${exclude_args[@]}" -f --dirsfirst --noreport
  else
    for dir in "${INCLUDE_DIRS[@]}"; do
      [ -d "$dir" ] && {
        printf "%b%s%b\n" "$y" "$dir" "$r"
        tree "${depth_arg[@]}" "${exclude_args[@]}" -f --dirsfirst --noreport "$dir"
        printf "\n"
      }
    done
  fi
  printf "\n"
}

generate_file_content() {
  local file_pattern_args=() dir_exclude_args=() file_exclude_args=()
  local rb=$RED_BOLD gr=$GRAY r=$RESET
  $USE_COLORS || rb="" gr="" r=""
  
  [ -n "$OUTPUT_FILE" ] && file_exclude_args+=("-not" "-path" "*/$OUTPUT_FILE" "-not" "-path" "./$OUTPUT_FILE")

  # Safely construct file pattern arguments
  if [ ${#INCLUDE_PATTERNS[@]} -gt 0 ]; then
    for pattern in "${INCLUDE_PATTERNS[@]}"; do
      if [ ${#file_pattern_args[@]} -eq 0 ]; then
        file_pattern_args+=("-name" "$pattern")
      else
        file_pattern_args+=("-o" "-name" "$pattern")
      fi
    done
  else
    # Default to all files if no patterns specified
    file_pattern_args=("-name" "*")
  fi
  for dir in "${EXCLUDE_DIRS[@]}"; do 
    # Match multiple possible path formats
    dir_exclude_args+=("-not" "-path" "*/$dir/*")
    dir_exclude_args+=("-not" "-path" "./$dir/*")
    dir_exclude_args+=("-not" "-path" "$dir/*")
    # Also exclude the directory itself
    dir_exclude_args+=("-not" "-path" "*/$dir")
    dir_exclude_args+=("-not" "-path" "./$dir")
    dir_exclude_args+=("-not" "-path" "$dir")
  done
  for file in "${EXCLUDE_FILES[@]}"; do file_exclude_args+=("-not" "-name" "$file"); done

  local search_dirs=(.) files=()
  [ ${#INCLUDE_DIRS[@]} -gt 0 ] && search_dirs=("${INCLUDE_DIRS[@]}")
  for dir in "${search_dirs[@]}"; do
    if [ -d "$dir" ]; then
      local depth_arg=()
      [ "$MAX_DEPTH" -ge 0 ] && depth_arg=("-maxdepth" "$MAX_DEPTH")
      while IFS= read -r f; do files+=("$f"); done < <(find "$dir" "${depth_arg[@]}" -type f "${file_pattern_args[@]}" "${dir_exclude_args[@]}" "${file_exclude_args[@]}" | sort)
    fi
  done

  # Additional filtering for excluded directories
  if [ ${#EXCLUDE_DIRS[@]} -gt 0 ]; then
    filtered_files=()
    for f in "${files[@]}"; do
      exclude=false
      for dir in "${EXCLUDE_DIRS[@]}"; do
        if [[ "$f" == *"/$dir/"* || "$f" == "./$dir/"* || "$f" == "$dir/"* ]]; then
          exclude=true
          break
        fi
      done
      $exclude || filtered_files+=("$f")
    done
    files=("${filtered_files[@]}")
  fi

  if [ -n "$SEARCH_PATTERN" ]; then
    local match=()
    for f in "${files[@]}"; do
      grep -q "$SEARCH_PATTERN" "$f" 2>/dev/null && match+=("$f")
    done
    files=("${match[@]}")
  fi

  case "$OUTPUT_FORMAT" in
  markdown)
    for f in "${files[@]}"; do
      printf "## %s\n\n" "$f"
      local ext="${f##*.}"
      printf '```%s\n' "$ext"
      if [ "$SHOW_LINE_NUMBERS" = true ]; then nl -ba "$f"; else cat "$f"; fi
      printf '```\n\n'
    done
    ;;
  json)
    echo "{"
    echo '  "files": ['
    local first=true
    for f in "${files[@]}"; do
      $first || echo "    },"
      first=false
      echo "    {"
      printf '      "path": "%s",\n' "$f"
      printf '      "content": %s\n' "$(python3 -c 'import json,sys;print(json.dumps(open(sys.argv[1]).read()))' "$f")"
    done
    [ ${#files[@]} -gt 0 ] && echo "    }"
    echo "  ]"
    echo "}"
    ;;
  *)
    for f in "${files[@]}"; do
      printf "%b**%s**%b\n" "$rb" "$f" "$r"
      if [ "$SHOW_LINE_NUMBERS" = true ]; then
        while IFS= read -r line; do
          printf "%b%s%b\n" "$gr" "$line" "$r"
        done < <(nl -ba "$f")
      else
        while IFS= read -r line; do
          printf "%b%s%b\n" "$gr" "$line" "$r"
        done <"$f"
      fi
      printf "\n"
    done
    ;;
  esac
}

# Parse args
while [[ $# -gt 0 ]]; do
  case $1 in
  -h | --help)
    print_usage
    exit 0
    ;;
  -i | --include)
    INCLUDE_PATTERNS+=("$2")
    shift 2
    ;;
  -e | --exclude-dir)
    EXCLUDE_DIRS+=("$2")
    shift 2
    ;;
  -x | --exclude-file)
    EXCLUDE_FILES+=("$2")
    shift 2
    ;;
  -d | --max-depth)
    MAX_DEPTH="$2"
    shift 2
    ;;
  -t | --no-tree)
    SHOW_TREE=false
    shift
    ;;
  -f | --no-files)
    SHOW_FILES=false
    shift
    ;;
  -n | --line-numbers)
    SHOW_LINE_NUMBERS=true
    shift
    ;;
  -o | --output)
    OUTPUT_FILE="$2"
    USE_COLORS=false  # Disable colors when output goes to a file
    shift 2
    ;;
  -s | --search)
    SEARCH_PATTERN="$2"
    shift 2
    ;;
  -p | --path)
    INCLUDE_DIRS+=("$2")
    shift 2
    ;;
  -m | --format)
    OUTPUT_FORMAT="$2"
    shift 2
    ;;
  *)
    local rb=$RED_BOLD r=$RESET
    $USE_COLORS || rb="" r=""
    printf "%bError: Unknown option %s%b\n" "$rb" "$1" "$r"
    print_usage
    exit 1
    ;;
  esac
done

# Remove excluded patterns
for excl in "${EXCLUDE_FILES[@]}"; do
  temp=()
  for pat in "${INCLUDE_PATTERNS[@]}"; do [[ "$pat" != "$excl" ]] && temp+=("$pat"); done
  INCLUDE_PATTERNS=("${temp[@]}")
done

# Validate format
[[ ! "$OUTPUT_FORMAT" =~ ^(text|markdown|json)$ ]] && {
  # Print the error message to stderr
  local rb=$RED_BOLD r=$RESET
  $USE_COLORS || rb="" r=""
  printf "%bError: Invalid output format. Must be one of: text, markdown, json%b\n" "$rb" "$r" >&2
  exit 1
}

# Redirect if needed
[ -n "$OUTPUT_FILE" ] && {
  TEMP_OUTPUT=$(mktemp)
  exec >"$TEMP_OUTPUT"
}

$SHOW_TREE && generate_tree
$SHOW_FILES && generate_file_content

[ -n "$TEMP_OUTPUT" ] && [ -n "$OUTPUT_FILE" ] && mv "$TEMP_OUTPUT" "$OUTPUT_FILE"

exit 0
