#!/bin/bash
set -e

LOKI_VERSION="${LOKI_VERSION:-3.0.0}"
SKILLS_DIR="${SKILLS_DIR:-/skills}"
LOKI_DATA_DIR="${LOKI_DATA_DIR:-/data/loki}"
SUPERVISOR_LOG_DIR="${SUPERVISOR_LOG_DIR:-/var/log/supervisor}"
SKILL_DIR="$(dirname "$(dirname "$0")")"

# Detect architecture
get_arch() {
    case "$(uname -m)" in
        x86_64)  echo "amd64" ;;
        aarch64) echo "arm64" ;;
        armv7l)  echo "arm" ;;
        *)
            echo "Unsupported architecture: $(uname -m)" >&2
            exit 1
            ;;
    esac
}

# Install Loki
install_loki() {
    if command -v loki &>/dev/null; then
        echo "loki already installed"
        return 0
    fi

    echo "Installing Loki v${LOKI_VERSION}..."
    local arch=$(get_arch)
    local url="https://github.com/grafana/loki/releases/download/v${LOKI_VERSION}/loki-linux-${arch}.zip"

    apt-get update && apt-get install -y unzip

    curl -sSL "$url" -o /tmp/loki.zip
    unzip -o /tmp/loki.zip -d /tmp
    mv /tmp/loki-linux-${arch} /usr/local/bin/loki
    chmod +x /usr/local/bin/loki
    rm /tmp/loki.zip

    echo "Loki installed"
}

# Install Promtail
install_promtail() {
    if command -v promtail &>/dev/null; then
        echo "promtail already installed"
        return 0
    fi

    echo "Installing Promtail v${LOKI_VERSION}..."
    local arch=$(get_arch)
    local url="https://github.com/grafana/loki/releases/download/v${LOKI_VERSION}/promtail-linux-${arch}.zip"

    curl -sSL "$url" -o /tmp/promtail.zip
    unzip -o /tmp/promtail.zip -d /tmp
    mv /tmp/promtail-linux-${arch} /usr/local/bin/promtail
    chmod +x /usr/local/bin/promtail
    rm /tmp/promtail.zip

    echo "Promtail installed"
}

# Setup directories
setup_dirs() {
    mkdir -p "$LOKI_DATA_DIR"
    mkdir -p "$SKILL_DIR/config"
    echo "Data directory: $LOKI_DATA_DIR"
}

# Generate Loki config
generate_loki_config() {
    local retention="${LOKI_RETENTION:-168h}"

    cat > "$SKILL_DIR/config/loki.yaml" << EOF
auth_enabled: false

server:
  http_listen_port: 3100
  grpc_listen_port: 9096

common:
  instance_addr: 127.0.0.1
  path_prefix: $LOKI_DATA_DIR
  storage:
    filesystem:
      chunks_directory: $LOKI_DATA_DIR/chunks
      rules_directory: $LOKI_DATA_DIR/rules
  replication_factor: 1
  ring:
    kvstore:
      store: inmemory

query_range:
  results_cache:
    cache:
      embedded_cache:
        enabled: true
        max_size_mb: 100

schema_config:
  configs:
    - from: 2020-10-24
      store: tsdb
      object_store: filesystem
      schema: v13
      index:
        prefix: index_
        period: 24h

limits_config:
  retention_period: $retention

ruler:
  alertmanager_url: http://localhost:9093

analytics:
  reporting_enabled: false
EOF

    echo "Loki config written to $SKILL_DIR/config/loki.yaml"
}

# Generate Promtail config with auto-discovered log paths
generate_promtail_config() {
    cat > "$SKILL_DIR/config/promtail.yaml" << EOF
server:
  http_listen_port: 9080
  grpc_listen_port: 0

positions:
  filename: /tmp/positions.yaml

clients:
  - url: http://localhost:3100/loki/api/v1/push

scrape_configs:
  # Supervisor logs (all skills)
  - job_name: skills
    static_configs:
      - targets:
          - localhost
        labels:
          job: skills
          __path__: $SUPERVISOR_LOG_DIR/*.log

    pipeline_stages:
      # Extract skill name from filename
      - match:
          selector: '{job="skills"}'
          stages:
            - regex:
                source: filename
                expression: '.*/(?P<skill>[^/]+)\\.log$'
            - labels:
                skill:

  # Supervisor error logs
  - job_name: skills_errors
    static_configs:
      - targets:
          - localhost
        labels:
          job: skills
          level: error
          __path__: $SUPERVISOR_LOG_DIR/*.err

    pipeline_stages:
      - match:
          selector: '{job="skills"}'
          stages:
            - regex:
                source: filename
                expression: '.*/(?P<skill>[^/]+)\\.err$'
            - labels:
                skill:
EOF

    echo "Promtail config written to $SKILL_DIR/config/promtail.yaml"
}

# Configure Caddy if present
configure_caddy() {
    local caddy_dir="$SKILLS_DIR/caddy"
    local loki_port="${LOKI_PORT:-3100}"
    local loki_domain="${LOKI_DOMAIN:-}"

    if [ ! -d "$caddy_dir" ]; then
        echo "Caddy not found - Loki API on port $loki_port"
        return 0
    fi

    echo "Caddy detected - configuring reverse proxy..."
    mkdir -p "$caddy_dir/snippets.d"

    local snippet="$caddy_dir/snippets.d/loki.caddy"

    if [ -n "$loki_domain" ]; then
        cat > "$snippet" << EOF
# Auto-generated by loki skill
$loki_domain {
    reverse_proxy localhost:$loki_port
}
EOF
        echo "Caddy config: $loki_domain -> localhost:$loki_port"
    else
        cat > "$snippet" << EOF
# Auto-generated by loki skill
# Add to your site block:
# handle /loki/* {
#     reverse_proxy localhost:$loki_port
# }
EOF
        echo "Caddy snippet created (manual config needed)"
    fi
}

install_loki
install_promtail
setup_dirs
generate_loki_config
generate_promtail_config
configure_caddy

echo "Loki setup complete"