#!/bin/bash
# Fetch logs and manage containers in the mcpctl Portainer stack
# Usage: ./logs.sh [container_name] [--tail N] [--since TIME] [--follow]
#        ./logs.sh restart <container>    Restart a container
#        ./logs.sh pull <image>           Pull an image on the NAS
#        ./logs.sh exec <container> <cmd> Run a command in a container

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

PORTAINER_URL="${PORTAINER_URL:-http://10.0.0.194:9000}"
PORTAINER_USER="${PORTAINER_USER:-michal}"
ENDPOINT_ID="2"

TAIL_LINES=100
SINCE=""
FOLLOW=false
SHOW_ALL=false
CONTAINER_NAME=""
TIMESTAMPS=true
COMMAND=""
CMD_TARGET=""
EXEC_CMD=()

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
DIM='\033[2m'
NC='\033[0m'

log_info() { echo -e "${GREEN}[INFO]${NC} $1" >&2; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1" >&2; }
log_error() { echo -e "${RED}[ERROR]${NC} $1" >&2; }

get_password() {
    if [[ -n "$PORTAINER_PASSWORD" ]]; then
        echo "$PORTAINER_PASSWORD"
        return
    fi
    if [[ -f "$SCRIPT_DIR/.portainer_password" ]]; then
        cat "$SCRIPT_DIR/.portainer_password"
        return
    fi
    if [[ -f "$HOME/.portainer_password" ]]; then
        cat "$HOME/.portainer_password"
        return
    fi
    read -s -p "Enter Portainer password for $PORTAINER_USER: " password
    echo >&2
    echo "$password"
}

get_jwt_token() {
    local password="$1"
    local escaped_password
    escaped_password=$(printf '%s' "$password" | jq -Rs .)

    local response
    response=$(curl -s -X POST "$PORTAINER_URL/api/auth" \
        -H "Content-Type: application/json" \
        -d "{\"Username\":\"$PORTAINER_USER\",\"Password\":$escaped_password}")

    local token
    token=$(echo "$response" | jq -r '.jwt // empty')

    if [[ -z "$token" ]]; then
        log_error "Authentication failed: $(echo "$response" | jq -r '.message // "Unknown error"')"
        exit 1
    fi
    echo "$token"
}

parse_since() {
    local since="$1"
    local now
    now=$(date +%s)

    if [[ "$since" =~ ^([0-9]+)s$ ]]; then
        echo $(( now - ${BASH_REMATCH[1]} ))
    elif [[ "$since" =~ ^([0-9]+)m$ ]]; then
        echo $(( now - ${BASH_REMATCH[1]} * 60 ))
    elif [[ "$since" =~ ^([0-9]+)h$ ]]; then
        echo $(( now - ${BASH_REMATCH[1]} * 3600 ))
    elif [[ "$since" =~ ^([0-9]+)d$ ]]; then
        echo $(( now - ${BASH_REMATCH[1]} * 86400 ))
    else
        date -d "$since" +%s 2>/dev/null || echo "$since"
    fi
}

list_containers() {
    local token="$1"

    local containers
    containers=$(curl -s -X GET "$PORTAINER_URL/api/endpoints/$ENDPOINT_ID/docker/containers/json?all=true" \
        -H "Authorization: Bearer $token")

    echo ""
    printf "${CYAN}%-30s %-15s %-20s %s${NC}\n" "CONTAINER" "STATE" "STATUS" "IMAGE"
    printf "${DIM}%-30s %-15s %-20s %s${NC}\n" "─────────" "─────" "──────" "─────"

    echo "$containers" | jq -r '.[] | [
        (.Names[0] // "unnamed" | ltrimstr("/")),
        .State,
        .Status,
        (.Image | split(":")[0] | split("/")[-1])
    ] | @tsv' | sort | while IFS=$'\t' read -r name state status image; do
        local color="$NC"
        case "$state" in
            running) color="$GREEN" ;;
            exited|dead) color="$RED" ;;
            restarting|paused) color="$YELLOW" ;;
        esac
        printf "%-30s ${color}%-15s${NC} %-20s %s\n" "$name" "$state" "$status" "$image"
    done

    echo ""
}

resolve_container() {
    local token="$1"
    local container_name="$2"

    local containers
    containers=$(curl -s -X GET "$PORTAINER_URL/api/endpoints/$ENDPOINT_ID/docker/containers/json?all=true" \
        -H "Authorization: Bearer $token")

    local container_id
    container_id=$(echo "$containers" | jq -r --arg name "$container_name" \
        '.[] | select(.Names[] | ltrimstr("/") == $name) | .Id')

    if [[ -z "$container_id" ]]; then
        container_id=$(echo "$containers" | jq -r --arg name "$container_name" \
            '.[] | select(.Names[] | ltrimstr("/") | contains($name)) | .Id' | head -1)
    fi

    if [[ -z "$container_id" ]]; then
        log_error "Container '$container_name' not found"
        echo "$containers" | jq -r '.[].Names[] | ltrimstr("/")' | sort >&2
        exit 1
    fi

    echo "$container_id"
}

fetch_logs() {
    local token="$1"
    local container_name="$2"

    local container_id
    container_id=$(resolve_container "$token" "$container_name")

    local query="stderr=true&stdout=true&tail=${TAIL_LINES}"
    if [[ "$TIMESTAMPS" == "true" ]]; then
        query="${query}&timestamps=true"
    fi
    if [[ -n "$SINCE" ]]; then
        local since_ts
        since_ts=$(parse_since "$SINCE")
        query="${query}&since=${since_ts}"
    fi

    if [[ "$FOLLOW" == "true" ]]; then
        query="${query}&follow=true"
        log_info "Streaming logs for $container_name (Ctrl+C to stop)..."
        curl -s -N -X GET "$PORTAINER_URL/api/endpoints/$ENDPOINT_ID/docker/containers/$container_id/logs?${query}" \
            -H "Authorization: Bearer $token" | \
            perl -pe 's/^.{8}//s'
    else
        curl -s -X GET "$PORTAINER_URL/api/endpoints/$ENDPOINT_ID/docker/containers/$container_id/logs?${query}" \
            -H "Authorization: Bearer $token" | \
            perl -pe 's/^.{8}//s'
    fi
}

fetch_all_logs() {
    local token="$1"

    local containers
    containers=$(curl -s -X GET "$PORTAINER_URL/api/endpoints/$ENDPOINT_ID/docker/containers/json?all=true" \
        -H "Authorization: Bearer $token")

    local names
    names=$(echo "$containers" | jq -r '.[].Names[] | ltrimstr("/")' | sort)

    if [[ "$TAIL_LINES" == "100" ]]; then
        TAIL_LINES=50
    fi

    while IFS= read -r name; do
        echo ""
        echo -e "${CYAN}=== ${name} ===${NC}"
        FOLLOW=false fetch_logs "$token" "$name" 2>/dev/null || echo -e "${DIM}  (no logs available)${NC}"
    done <<< "$names"
}

restart_container() {
    local token="$1"
    local container_name="$2"

    local container_id
    container_id=$(resolve_container "$token" "$container_name")

    log_info "Restarting $container_name..."
    local http_code
    http_code=$(curl -s -o /dev/null -w "%{http_code}" -X POST \
        "$PORTAINER_URL/api/endpoints/$ENDPOINT_ID/docker/containers/$container_id/restart" \
        -H "Authorization: Bearer $token")

    if [[ "$http_code" == "204" ]]; then
        log_info "$container_name restarted successfully"
    else
        log_error "Restart failed (HTTP $http_code)"
        exit 1
    fi
}

pull_image() {
    local token="$1"
    local image="$2"

    local image_name="${image%%:*}"
    local tag="${image#*:}"
    [[ "$tag" == "$image" ]] && tag="latest"

    log_info "Pulling $image_name:$tag on NAS..."
    local response
    response=$(curl -s -X POST \
        "$PORTAINER_URL/api/endpoints/$ENDPOINT_ID/docker/images/create?fromImage=$image_name&tag=$tag" \
        -H "Authorization: Bearer $token")

    if echo "$response" | grep -q '"error"'; then
        log_error "Pull failed:"
        echo "$response" | grep '"error"'
        exit 1
    fi

    local status
    status=$(echo "$response" | tail -1 | jq -r '.status // empty' 2>/dev/null)
    log_info "$status"
}

exec_container() {
    local token="$1"
    local container_name="$2"
    shift 2
    local cmd=("$@")

    local container_id
    container_id=$(resolve_container "$token" "$container_name")

    local cmd_json
    cmd_json=$(printf '%s\n' "${cmd[@]}" | jq -R . | jq -s '.')

    local exec_id
    exec_id=$(curl -s -X POST \
        "$PORTAINER_URL/api/endpoints/$ENDPOINT_ID/docker/containers/$container_id/exec" \
        -H "Authorization: Bearer $token" \
        -H "Content-Type: application/json" \
        -d "{\"AttachStdout\":true,\"AttachStderr\":true,\"Cmd\":$cmd_json}" | jq -r '.Id')

    if [[ -z "$exec_id" || "$exec_id" == "null" ]]; then
        log_error "Failed to create exec instance"
        exit 1
    fi

    curl -s -X POST \
        "$PORTAINER_URL/api/endpoints/$ENDPOINT_ID/docker/exec/$exec_id/start" \
        -H "Authorization: Bearer $token" \
        -H "Content-Type: application/json" \
        -d '{"Detach":false,"Tty":true}'
}

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --tail|-n)
            TAIL_LINES="$2"
            shift 2
            ;;
        --since)
            SINCE="$2"
            shift 2
            ;;
        --follow|-f)
            FOLLOW=true
            shift
            ;;
        --all|-a)
            SHOW_ALL=true
            shift
            ;;
        --no-timestamps)
            TIMESTAMPS=false
            shift
            ;;
        --help|-h)
            echo "Usage: $0 [container_name] [options]"
            echo "       $0 restart <container>"
            echo "       $0 pull <image>"
            echo "       $0 exec <container> <command...>"
            echo ""
            echo "Fetch logs and manage containers in the mcpctl Portainer stack."
            echo ""
            echo "LOGS:"
            echo "  (no args)                   List all containers with status"
            echo "  <container>                 Last 100 lines from container"
            echo "  --all, -a                   Show logs from all containers"
            echo "  --tail N, -n N              Number of lines (default: 100)"
            echo "  --since TIME                Logs since TIME (e.g., 1h, 30m, 2d)"
            echo "  --follow, -f                Stream logs in real-time"
            echo "  --no-timestamps             Hide timestamps"
            echo ""
            echo "MANAGEMENT:"
            echo "  restart <container>         Restart a container"
            echo "  pull <image[:tag]>          Pull an image on the NAS"
            echo "  exec <container> <cmd...>   Run a command in a container"
            exit 0
            ;;
        restart)
            COMMAND="restart"
            CMD_TARGET="${2:?Usage: $0 restart <container>}"
            shift 2
            ;;
        pull)
            COMMAND="pull"
            CMD_TARGET="${2:?Usage: $0 pull <image[:tag]>}"
            shift 2
            ;;
        exec)
            COMMAND="exec"
            CMD_TARGET="${2:?Usage: $0 exec <container> <command...>}"
            shift 2
            EXEC_CMD=("$@")
            if [[ ${#EXEC_CMD[@]} -eq 0 ]]; then
                log_error "Usage: $0 exec <container> <command...>"
                exit 1
            fi
            break
            ;;
        -*)
            log_error "Unknown option: $1"
            exit 1
            ;;
        *)
            CONTAINER_NAME="$1"
            shift
            ;;
    esac
done

# Main
password=$(get_password)
token=$(get_jwt_token "$password")

if [[ "$COMMAND" == "restart" ]]; then
    restart_container "$token" "$CMD_TARGET"
elif [[ "$COMMAND" == "pull" ]]; then
    pull_image "$token" "$CMD_TARGET"
elif [[ "$COMMAND" == "exec" ]]; then
    exec_container "$token" "$CMD_TARGET" "${EXEC_CMD[@]}"
elif [[ "$SHOW_ALL" == "true" ]]; then
    fetch_all_logs "$token"
elif [[ -n "$CONTAINER_NAME" ]]; then
    fetch_logs "$token" "$CONTAINER_NAME"
else
    list_containers "$token"
fi