#!/usr/bin/env bash
#
# Imajin CLI Wrapper
# Routes to either service CLI or test runner
#

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CLI_SCRIPT="$SCRIPT_DIR/packages/cli/imajin"
VENV_PYTHON="$SCRIPT_DIR/orchestrators/imajin-pipeline/.venv/bin/python"
PIPELINE_DIR="$SCRIPT_DIR/orchestrators/imajin-pipeline"
APP_DIR="$SCRIPT_DIR/orchestrators/imajin-app"

# Show help if no args
if [[ $# -eq 0 ]]; then
    echo "Usage: ./run <command> [args...]"
    echo ""
    echo "Test commands:"
    echo "  test [target]              Run tests (unit by default)"
    echo "  test unit                  Run unit tests (fast, mocked)"
    echo "  test contracts             Run contract/type unit tests (batch, strategy)"
    echo "  test integration           Run internal pipeline integration tests (GPU)"
    echo "  test orchestrator          Run HTTP orchestrator E2E tests (port 8080, GPU)"
    echo "  test orchestrator --save-images   Save generated images to /tmp/"
    echo "  test <pattern>             Run tests matching pattern"
    echo "  test -k 'expr'             Run tests matching pytest expression"
    echo ""
    echo "Service commands (imajin services):"
    echo "  start [service]       Start all services or specific service"
    echo "  stop [service]        Stop all services or specific service"
    echo "  health                Check health of all services"
    echo "  logs <service>        Tail logs for service"
    echo ""
    echo "Examples:"
    echo "  ./run test                               # Run unit tests"
    echo "  ./run test integration                   # Run GPU integration tests"
    echo "  ./run test orchestrator                  # Run HTTP orchestrator E2E tests"
    echo "  ./run test orchestrator --save-images   # Save generated images to /tmp/"
    echo "  ./run test test_outfit_parser            # Run specific test"
    echo "  ./run start diffusion                    # Start diffusion service"
    exit 0
fi

# Check if venv exists
if [[ ! -f "$VENV_PYTHON" ]]; then
    echo "❌ Orchestrator venv not found at $VENV_PYTHON"
    echo "   Run: cd orchestrators/imajin-pipeline && python3 -m venv .venv && pip install -e '.[dev]'"
    exit 1
fi

# Route based on first argument
case "$1" in
    test|tests)
        shift

        # Extract --output=<dir> and other flags from args
        OUTPUT_DIR=""
        EXTRA_ARGS=()
        TARGET="unit"
        TARGET_SET=false

        for arg in "$@"; do
            if [[ "$arg" == --output=* ]]; then
                OUTPUT_DIR="${arg#--output=}"
            elif [[ "$TARGET_SET" == false && "$arg" != -* ]]; then
                TARGET="$arg"
                TARGET_SET=true
            else
                EXTRA_ARGS+=("$arg")
            fi
        done

        # Build pytest command with PYTHONPATH set to pipeline src
        CMD=(env PYTHONPATH="$PIPELINE_DIR/src" "$VENV_PYTHON" -m pytest)

        case "$TARGET" in
            unit)
                CMD+=("$PIPELINE_DIR/tests/unit")
                ;;
            integration)
                CMD+=("$PIPELINE_DIR/tests/integration" --gpu)
                ;;
            contracts)
                # Contract/type unit tests (batch contracts, strategy, types)
                CMD=(env PYTHONPATH="$APP_DIR/src" "$VENV_PYTHON" -m pytest)
                CMD+=("$APP_DIR/tests" --ignore="$APP_DIR/tests/integration")
                ;;
            orchestrator)
                # HTTP orchestrator E2E tests (imajin-app/tests/integration/)
                CMD=(env PYTHONPATH="$APP_DIR/src" "$VENV_PYTHON" -m pytest)
                CMD+=("$APP_DIR/tests/integration" --gpu)
                ;;
            all)
                CMD+=("$PIPELINE_DIR/tests")
                [[ "${EXTRA_ARGS[*]}" == *"--gpu"* ]] || CMD+=(--gpu)
                ;;
            *)
                # Treat as test name pattern
                CMD+=("$PIPELINE_DIR/tests" -k "$TARGET")
                ;;
        esac

        # Add output directory if specified
        if [[ -n "$OUTPUT_DIR" ]]; then
            CMD+=(--output="$OUTPUT_DIR")
        fi

        # Add remaining args (-v, --gpu, etc.)
        CMD+=("${EXTRA_ARGS[@]}")

        # Add default flags
        [[ "${EXTRA_ARGS[*]}" == *"--tb"* ]] || CMD+=(--tb=short)

        echo "Running: ${CMD[*]}"
        echo ""
        exec "${CMD[@]}"
        ;;

    start|stop|health|logs)
        # Route to service CLI
        if [[ ! -f "$CLI_SCRIPT" ]]; then
            echo "❌ CLI not found at $CLI_SCRIPT"
            exit 1
        fi
        chmod +x "$CLI_SCRIPT"
        exec "$VENV_PYTHON" "$CLI_SCRIPT" "$@"
        ;;

    *)
        # Unknown command - try service CLI
        if [[ -f "$CLI_SCRIPT" ]]; then
            chmod +x "$CLI_SCRIPT"
            exec "$VENV_PYTHON" "$CLI_SCRIPT" "$@"
        else
            echo "❌ Unknown command: $1"
            echo "   Run './run' for help"
            exit 1
        fi
        ;;
esac