#!/usr/bin/env bash
# ampviz - amp execute with visualization
# Usage: ampviz "prompt" or echo "prompt" | ampviz

# Resolve symlink to get actual script location
SOURCE="${BASH_SOURCE[0]}"
while [ -L "$SOURCE" ]; do
  DIR="$(cd -P "$(dirname "$SOURCE")" && pwd)"
  SOURCE="$(readlink "$SOURCE")"
  [[ $SOURCE != /* ]] && SOURCE="$DIR/$SOURCE"
done
SCRIPT_DIR="$(cd -P "$(dirname "$SOURCE")/.." && pwd)"

if [ -t 0 ] && [ $# -gt 0 ]; then
  # Args provided, no stdin
  amp --execute "$@" --stream-json --dangerously-allow-all |
    npx --prefix "$SCRIPT_DIR" tsx -e "import { visualize } from '$SCRIPT_DIR/src/commands/visualize'; visualize();"
elif [ ! -t 0 ]; then
  # Stdin provided
  if [ $# -gt 0 ]; then
    # Both stdin and args
    cat | amp --execute "$@" --stream-json --dangerously-allow-all |
      npx --prefix "$SCRIPT_DIR" tsx -e "import { visualize } from '$SCRIPT_DIR/src/commands/visualize'; visualize();"
  else
    # Just stdin
    cat | amp --execute --stream-json --dangerously-allow-all |
      npx --prefix "$SCRIPT_DIR" tsx -e "import { visualize } from '$SCRIPT_DIR/src/commands/visualize'; visualize();"
  fi
else
  echo "Usage: ampviz \"prompt\""
  echo "       echo \"prompt\" | ampviz"
  echo "       cat file.md | ampviz"
  exit 1
fi