#!/bin/bash
#
# Dune Weaver CLI
#
# Usage: dw <command>
#
# Commands:
#   install     Run initial setup (Docker + WiFi fix)
#   start       Start Dune Weaver
#   stop        Stop Dune Weaver
#   restart     Restart Dune Weaver
#   update      Pull latest changes and restart
#   logs        View live logs (Ctrl+C to exit)
#   status      Show container status
#   shell       Open a shell in the container
#   checkout    Switch to a branch and pull its Docker images
#   touch       Manage touch screen app
#   help        Show this help message
#

set -e

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

# Find dune-weaver directory
find_install_dir() {
    # Check common locations
    if [[ -f "$HOME/dune-weaver/main.py" ]]; then
        echo "$HOME/dune-weaver"
    elif [[ -f "/home/pi/dune-weaver/main.py" ]]; then
        echo "/home/pi/dune-weaver"
    elif [[ -f "./main.py" ]]; then
        pwd
    else
        echo ""
    fi
}

# Check if using Docker or systemd
is_docker_mode() {
    # Docker mode if docker-compose.yml exists and docker is available
    [[ -f "$INSTALL_DIR/docker-compose.yml" ]] && command -v docker &> /dev/null
}

# Set IMAGE_TAG from current git branch (e.g. feature/foo -> feature-foo)
set_image_tag() {
    local branch
    branch=$(git -C "$INSTALL_DIR" rev-parse --abbrev-ref HEAD)
    export IMAGE_TAG="${branch//\//-}"
}

INSTALL_DIR=$(find_install_dir)

# Check if installed
check_installed() {
    if [[ -z "$INSTALL_DIR" ]]; then
        echo -e "${RED}Error: Dune Weaver not found${NC}"
        echo ""
        echo "Please install first:"
        echo "  git clone https://github.com/tuanchris/dune-weaver --single-branch"
        echo "  cd dune-weaver"
        echo "  dw install"
        exit 1
    fi
}

# Commands
cmd_install() {
    if [[ -z "$INSTALL_DIR" ]]; then
        # Not installed, check if we're in the right directory
        if [[ -f "./docker-compose.yml" ]] && [[ -f "./main.py" ]]; then
            INSTALL_DIR=$(pwd)
        else
            echo -e "${RED}Error: Run this from the dune-weaver directory${NC}"
            echo ""
            echo "  git clone https://github.com/tuanchris/dune-weaver --single-branch"
            echo "  cd dune-weaver"
            echo "  dw install"
            exit 1
        fi
    fi

    cd "$INSTALL_DIR"
    bash setup-pi.sh "$@"
}

cmd_start() {
    check_installed
    echo -e "${BLUE}Starting Dune Weaver...${NC}"
    cd "$INSTALL_DIR"

    if is_docker_mode; then
        set_image_tag
        sudo -E docker compose up -d
    else
        sudo systemctl start dune-weaver
    fi

    local fe_port="${FRONTEND_PORT:-80}"
    local port_suffix=""
    [[ "$fe_port" != "80" ]] && port_suffix=":$fe_port"
    echo -e "${GREEN}Started!${NC} Access at http://$(hostname -I | awk '{print $1}')${port_suffix}"
}

cmd_stop() {
    check_installed
    echo -e "${BLUE}Stopping Dune Weaver...${NC}"
    cd "$INSTALL_DIR"

    if is_docker_mode; then
        sudo docker compose down
    else
        sudo systemctl stop dune-weaver
    fi

    echo -e "${GREEN}Stopped${NC}"
}

cmd_restart() {
    check_installed
    echo -e "${BLUE}Restarting Dune Weaver...${NC}"
    cd "$INSTALL_DIR"

    if is_docker_mode; then
        set_image_tag
        sudo -E docker compose restart
    else
        sudo systemctl restart dune-weaver
    fi

    echo -e "${GREEN}Restarted${NC}"
}

cmd_update() {
    check_installed
    echo -e "${BLUE}Updating Dune Weaver...${NC}"
    cd "$INSTALL_DIR"

    # Check if we should skip the pull phase (called after re-exec)
    if [[ "$1" != "--continue" ]]; then
        echo "Pulling latest code..."
        git pull

        # Update dw CLI
        echo "Updating dw command..."
        sudo cp "$INSTALL_DIR/dw" /usr/local/bin/dw
        sudo chmod +x /usr/local/bin/dw

        # Re-exec with the new script to ensure new code runs
        echo "Restarting with updated CLI..."
        exec /usr/local/bin/dw update --continue
    fi

    if is_docker_mode; then
        set_image_tag
        echo -e "Image tag: ${BLUE}${IMAGE_TAG}${NC}"

        echo "Pulling latest Docker image..."
        sudo -E docker compose pull

        echo "Stopping current container..."
        sudo docker compose down

        echo "Starting with new version..."
        sudo -E docker compose up -d --remove-orphans

        echo "Cleaning up unused Docker resources..."
        sudo docker image prune -f
        sudo docker container prune -f
    else
        echo "Updating Python dependencies..."
        source .venv/bin/activate
        pip install -r requirements.txt

        echo "Restarting service..."
        sudo systemctl restart dune-weaver
    fi

    # Update touch app if installed (check both directory and service exist)
    local touch_dir="$INSTALL_DIR/dune-weaver-touch"
    if [[ -d "$touch_dir" ]] && systemctl list-unit-files dune-weaver-touch.service 2>/dev/null | grep -q dune-weaver-touch; then
        echo ""
        echo -e "${BLUE}Updating touch app...${NC}"
        cd "$touch_dir"

        # Update Python dependencies (code already pulled with main repo)
        if [[ -f "requirements.txt" ]] && [[ -d ".venv" ]]; then
            echo "Updating touch app dependencies..."
            source .venv/bin/activate
            pip install -q -r requirements.txt
            deactivate
        fi

        # Restart to apply changes
        sudo systemctl restart dune-weaver-touch
        echo -e "${GREEN}Touch app updated!${NC}"
    fi

    echo -e "${GREEN}Update complete!${NC}"
}

cmd_logs() {
    check_installed
    cd "$INSTALL_DIR"

    # Parse optional line count (e.g., dw logs 100)
    local lines="${1:-}"
    local follow="-f"

    if [[ -n "$lines" ]]; then
        follow="--tail $lines"
    fi

    if is_docker_mode; then
        sudo docker compose logs $follow
    else
        if [[ -n "$lines" ]]; then
            sudo journalctl -u dune-weaver -n "$lines"
        else
            sudo journalctl -u dune-weaver -f
        fi
    fi
}

cmd_status() {
    check_installed
    cd "$INSTALL_DIR"

    if is_docker_mode; then
        sudo docker compose ps
    else
        sudo systemctl status dune-weaver
    fi
}

cmd_shell() {
    check_installed
    cd "$INSTALL_DIR"

    if is_docker_mode; then
        sudo docker compose exec dune-weaver /bin/bash
    else
        echo -e "${YELLOW}Shell not available in systemd mode${NC}"
        echo "Use: cd $INSTALL_DIR && source .venv/bin/activate"
    fi
}

# Touch app commands (systemd service)
cmd_touch() {
    local subcmd="${1:-help}"

    case "$subcmd" in
        logs)
            local lines="${2:-}"
            if [[ -n "$lines" ]]; then
                sudo journalctl -u dune-weaver-touch -n "$lines"
            else
                sudo journalctl -u dune-weaver-touch -f
            fi
            ;;
        restart)
            echo -e "${BLUE}Restarting touch app...${NC}"
            sudo systemctl restart dune-weaver-touch
            echo -e "${GREEN}Touch app restarted${NC}"
            ;;
        stop)
            echo -e "${BLUE}Stopping touch app...${NC}"
            sudo systemctl stop dune-weaver-touch
            echo -e "${GREEN}Touch app stopped${NC}"
            ;;
        start)
            echo -e "${BLUE}Starting touch app...${NC}"
            sudo systemctl start dune-weaver-touch
            echo -e "${GREEN}Touch app started${NC}"
            ;;
        status)
            sudo systemctl status dune-weaver-touch
            ;;
        install)
            check_installed
            local touch_dir="$INSTALL_DIR/dune-weaver-touch"

            if [[ ! -f "$touch_dir/install.sh" ]]; then
                echo -e "${RED}Touch app installer not found at $touch_dir/install.sh${NC}"
                echo "Make sure you have the dune-weaver-touch directory in your dune-weaver folder."
                exit 1
            fi

            echo -e "${BLUE}Running touch app installer...${NC}"
            cd "$touch_dir"
            sudo bash install.sh
            ;;
        help|*)
            echo -e "${GREEN}Touch App Commands${NC}"
            echo ""
            echo "Usage: dw touch <command>"
            echo ""
            echo "Commands:"
            echo "  install     Install touch app (clone, setup, enable service)"
            echo "  logs [N]    View logs (N lines or follow if omitted)"
            echo "  restart     Restart touch app"
            echo "  stop        Stop touch app"
            echo "  start       Start touch app"
            echo "  status      Show touch app status"
            echo "  help        Show this help message"
            ;;
    esac
}

cmd_checkout() {
    check_installed
    local branch="$1"

    if [[ -z "$branch" ]]; then
        echo -e "${RED}Usage: dw checkout <branch>${NC}"
        echo ""
        echo "Examples:"
        echo "  dw checkout main"
        echo "  dw checkout feature/security-settings-split"
        echo ""
        echo "Current branch: $(git -C "$INSTALL_DIR" rev-parse --abbrev-ref HEAD)"
        exit 1
    fi

    cd "$INSTALL_DIR"

    # Repo may be cloned with --single-branch; fetch the target branch explicitly
    echo -e "${BLUE}Fetching branch ${branch}...${NC}"
    if ! git fetch origin "$branch" 2>/dev/null; then
        echo -e "${RED}Branch '${branch}' not found on remote${NC}"
        exit 1
    fi

    # Checkout: create local tracking branch if it doesn't exist
    if git show-ref --verify --quiet "refs/heads/$branch"; then
        git checkout "$branch"
        git pull
    else
        git checkout -b "$branch" "origin/$branch"
    fi

    echo -e "Switched to branch: ${GREEN}${branch}${NC}"

    # Pull and restart Docker images for the new branch
    if is_docker_mode; then
        set_image_tag
        echo -e "Image tag: ${BLUE}${IMAGE_TAG}${NC}"

        echo "Pulling Docker images..."
        sudo -E docker compose pull

        echo "Restarting with new images..."
        sudo docker compose down
        sudo -E docker compose up -d --remove-orphans

        echo "Cleaning up old images..."
        sudo docker image prune -f
    fi

    # Update dw CLI from the new branch
    sudo cp "$INSTALL_DIR/dw" /usr/local/bin/dw
    sudo chmod +x /usr/local/bin/dw

    echo -e "${GREEN}Checkout complete!${NC}"
}

cmd_help() {
    echo -e "${GREEN}Dune Weaver CLI${NC}"
    echo ""
    echo "Usage: dw <command>"
    echo ""
    echo "Commands:"
    echo "  install     Run initial setup (Docker + WiFi fix)"
    echo "  start       Start Dune Weaver"
    echo "  stop        Stop Dune Weaver"
    echo "  restart     Restart Dune Weaver"
    echo "  update      Pull latest changes and restart"
    echo "  checkout    Switch to a branch and pull its Docker images"
    echo "  logs [N]    View logs (N=number of lines, or follow if omitted)"
    echo "  status      Show container status"
    echo "  shell       Open a shell in the container"
    echo "  touch       Manage touch screen app (run 'dw touch help')"
    echo "  help        Show this help message"
    echo ""
    if [[ -n "$INSTALL_DIR" ]]; then
        echo -e "Install directory: ${BLUE}$INSTALL_DIR${NC}"
    fi
}

# Main
case "${1:-help}" in
    install)
        shift
        cmd_install "$@"
        ;;
    start)
        cmd_start
        ;;
    stop)
        cmd_stop
        ;;
    restart)
        cmd_restart
        ;;
    update)
        cmd_update "$2"
        ;;
    checkout)
        cmd_checkout "$2"
        ;;
    logs)
        cmd_logs "$2"
        ;;
    status)
        cmd_status
        ;;
    shell)
        cmd_shell
        ;;
    touch)
        shift
        cmd_touch "$@"
        ;;
    help|--help|-h)
        cmd_help
        ;;
    *)
        echo -e "${RED}Unknown command: $1${NC}"
        echo ""
        cmd_help
        exit 1
        ;;
esac