#!/usr/bin/env bash

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Check if pyright is installed
if ! command -v pyright &> /dev/null; then
    echo -e "${RED}Error: pyright is not installed${NC}"
    exit 1
fi

# Function to display usage
usage() {
    echo "Usage:"
    echo "  bash run_pyright.sh file1.py file2.py ..."
    echo "  bash run_pyright.sh --dir /path/to/directory"
    exit 1
}

# Check if arguments are provided
if [[ $# -eq 0 ]]; then
    usage
fi

# Parse arguments
files=()
if [[ "$1" == "--dir" ]]; then
    if [[ -z "${2:-}" ]]; then
        echo -e "${RED}Error: --dir requires a path argument${NC}"
        exit 1
    fi
    if [[ ! -d "$2" ]]; then
        echo -e "${RED}Error: Directory '$2' does not exist${NC}"
        exit 1
    fi
    # Find all .py files in the directory
    while IFS= read -r file; do
        files+=("$file")
    done < <(find "$2" -type f -name "*.py" | sort)

    if [[ ${#files[@]} -eq 0 ]]; then
        echo -e "${YELLOW}Warning: No Python files found in '$2'${NC}"
        exit 0
    fi
else
    # Use all arguments as file paths
    files=("$@")
fi

# Run pyright
echo -e "${GREEN}Running pyright on ${#files[@]} file(s)...${NC}"
echo ""

# Find project root (directory containing pyrightconfig.json) from the first file
project_root=""
if [[ ${#files[@]} -gt 0 ]]; then
    first_file="${files[0]}"
    search_dir="$(dirname "$first_file")"
    while [[ "$search_dir" != "/" ]]; do
        if [[ -f "$search_dir/pyrightconfig.json" ]]; then
            project_root="$search_dir"
            break
        fi
        search_dir="$(dirname "$search_dir")"
    done
fi

if [[ -n "$project_root" ]]; then
    cd "$project_root" || true
fi

timeout 120 pyright "${files[@]}"
exit_code=$?
if [[ $exit_code -eq 124 ]]; then
    echo -e "${YELLOW}⚠ pyright timed out after 120 seconds${NC}"
fi

# Summary
echo ""
if [[ $exit_code -eq 0 ]]; then
    echo -e "${GREEN}✓ All files passed type checking${NC}"
else
    echo -e "${RED}✗ Type checking found issues (exit code: $exit_code)${NC}"
fi

exit $exit_code