#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Installer & updater for Grafana Loki (Debian 13 / trixie).
Default run uses:
  --deb-url latest from GitHub
  --config /etc/loki/custom-config.yml
  --data-dir /data/loki

--force: wykonuje wszystkie kroki od nowa i nadpisuje pliki/katalogi/konfigi.
"""

import argparse
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
import urllib.request
from urllib.error import URLError, HTTPError

GITHUB_LATEST_API = "https://api.github.com/repos/grafana/loki/releases/latest"

PKG_NAME = "loki"
BIN_PATH = "/usr/bin/loki"

SYSTEMD_UNIT = "/lib/systemd/system/loki.service"
SYSTEMD_OVERRIDE_DIR = "/etc/systemd/system/loki.service.d"
SYSTEMD_OVERRIDE_PATH = os.path.join(SYSTEMD_OVERRIDE_DIR, "override.conf")

DEFAULT_USER = "loki"
DEFAULT_GROUP = "loki"
DEFAULT_CONFIG_DIR = "/etc/loki"
DEFAULT_CONFIG_PATH = os.path.join(DEFAULT_CONFIG_DIR, "custom-config.yml")
DEFAULT_DATA_DIR = "/data/loki"
DEFAULT_LOG_DIR = "/var/log/loki"

CUSTOM_CONFIG = """
auth_enabled: false

server:
  http_listen_port: 3100
  grpc_listen_port: 9096
  log_level: error
  grpc_server_max_concurrent_streams: 1000

common:
  instance_addr: 0.0.0.0
  path_prefix: /data/loki
  storage:
    filesystem:
      chunks_directory: /data/loki/chunks
      rules_directory: /data/loki/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

storage_config:
  filesystem:
    directory: /data/loki/db

pattern_ingester:
  enabled: true
  metric_aggregation:
    loki_address: 0.0.0.0:3100

ruler:
  alertmanager_url: http://10.87.2.22:9093

frontend:
  encoding: protobuf

query_scheduler:
  max_outstanding_requests_per_tenant: 2048

limits_config:
  ingestion_rate_mb: 1024
  ingestion_burst_size_mb: 1024
  max_query_lookback: 24h
  max_query_series: 10000
  retention_period: 672h

compactor:
  working_directory: /data/loki/compactor
  compaction_interval: 1h
  retention_enabled: true
  retention_delete_delay: 30m
  retention_delete_worker_count: 150
  delete_request_store: filesystem
"""

# ---------------- helpers ----------------

def log(msg: str, quiet: bool):
    if not quiet:
        print(msg)

def run(cmd, check=True, quiet=False):
    p = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if check and p.returncode != 0:
        if quiet:
            sys.stderr.write(p.stderr or p.stdout)
        else:
            print(p.stdout)
            print(p.stderr, file=sys.stderr)
        raise SystemExit(p.returncode)
    return p.returncode, p.stdout.strip(), p.stderr.strip()

def ensure_user(user=DEFAULT_USER, group=DEFAULT_GROUP, quiet=False):
    code, _, _ = run(f"getent group {group}", check=False, quiet=quiet)
    if code != 0:
        log(f"Creating group: {group}", quiet)
        run(f"groupadd --system {group}")
    code, _, _ = run(f"id -u {user}", check=False, quiet=quiet)
    if code != 0:
        log(f"Creating user: {user}", quiet)
        run(
            "useradd --system --no-create-home --shell /usr/sbin/nologin "
            f"--gid {group} {user}"
        )

def ensure_dirs(paths, owner=(DEFAULT_USER, DEFAULT_GROUP), mode=0o750, quiet=False, force=False):
    for p in paths:
        if force and os.path.isdir(p):
            shutil.rmtree(p)
        if not os.path.isdir(p):
            log(f"Creating dir: {p}", quiet)
            os.makedirs(p, exist_ok=True)
        os.chmod(p, mode)
        run(f"chown -R {owner[0]}:{owner[1]} {p}")

def write_config(path, content, mode=0o640, owner=(DEFAULT_USER, DEFAULT_GROUP), quiet=False, force=False):
    if force or not os.path.exists(path):
        base = os.path.dirname(path)
        if base and not os.path.isdir(base):
            os.makedirs(base, exist_ok=True)
        log(f"Writing config: {path}", quiet)
        with open(path, "w", encoding="utf-8") as f:
            f.write(content)
        os.chmod(path, mode)
        run(f"chown {owner[0]}:{owner[1]} {path}")

def download(url, dest, quiet=False):
    log(f"Downloading: {url}", quiet)
    req = urllib.request.Request(url, headers={"User-Agent": "loki-installer"})
    with urllib.request.urlopen(req) as r, open(dest, "wb") as f:
        shutil.copyfileobj(r, f)

def get_installed_version(quiet=False):
    code, out, _ = run(f"dpkg-query -W -f='${{Version}}\n' {PKG_NAME}", check=False, quiet=quiet)
    return out.strip() if code == 0 else None

def semver_tuple(v: str):
    m = re.match(r"(\d+)(?:\.(\d+))?(?:\.(\d+))?", v)
    if not m:
        return (0, 0, 0)
    parts = [int(x) if x else 0 for x in m.groups()]
    while len(parts) < 3:
        parts.append(0)
    return tuple(parts)

def find_latest_loki_deb_asset(quiet=False):
    req = urllib.request.Request(GITHUB_LATEST_API, headers={"User-Agent": "loki-installer", "Accept": "application/vnd.github+json"})
    with urllib.request.urlopen(req) as r:
        data = json.loads(r.read().decode("utf-8"))
    for a in data.get("assets", []):
        name = a.get("name", "")
        url = a.get("browser_download_url")
        if name.startswith("loki_") and name.endswith("_amd64.deb"):
            ver = name.split("_")[1]
            return ver, url
    return None, None

def install_deb_from_url(url, quiet=False, force=False):
    with tempfile.TemporaryDirectory() as td:
        deb_path = os.path.join(td, os.path.basename(url) or "loki.deb")
        download(url, deb_path, quiet=quiet)
        log("Installing package…", quiet)
        run("apt-get update", check=False, quiet=quiet)
        run(f"DEBIAN_FRONTEND=noninteractive apt-get -o Dpkg::Options::=\"--force-confdef\" -o Dpkg::Options::=\"--force-confold\" -y install {deb_path}")
        log("Installed loki", quiet)

def ensure_systemd(config_path=DEFAULT_CONFIG_PATH, user=DEFAULT_USER, group=DEFAULT_GROUP, quiet=False, force=False):
    unit_path = "/etc/systemd/system/loki.service" if not os.path.isfile(SYSTEMD_UNIT) else None
    if force and os.path.exists("/etc/systemd/system/loki.service"):
        os.remove("/etc/systemd/system/loki.service")
    if force and os.path.exists(SYSTEMD_OVERRIDE_PATH):
        os.remove(SYSTEMD_OVERRIDE_PATH)
    if unit_path:
        unit = f"""
[Unit]
Description=Grafana Loki
After=network.target

[Service]
Type=simple
ExecStart={BIN_PATH} -config.file={config_path} #-target=compactor
User={user}
Group={group}
Restart=always
RestartSec=5

[Install]
WantedBy=multi-user.target
"""
        with open(unit_path, "w", encoding="utf-8") as f:
            f.write(unit.lstrip())
    else:
        os.makedirs(SYSTEMD_OVERRIDE_DIR, exist_ok=True)
        override = f"""
[Service]
User={user}
Group={group}
ExecStart=
ExecStart={BIN_PATH} -config.file={config_path}
"""
        with open(SYSTEMD_OVERRIDE_PATH, "w", encoding="utf-8") as f:
            f.write(override.lstrip())

    run("systemctl daemon-reload")
    run("systemctl enable --now loki")

def main():
    p = argparse.ArgumentParser(description="Install/Update Grafana Loki (Debian 13)")
    p.add_argument("--deb-url", default="latest", help="URL to loki amd64 .deb or 'latest'")
    p.add_argument("--update", action="store_true", help="Silent update from GitHub")
    p.add_argument("--force", action="store_true", help="Redo all steps, overwrite config and dirs")
    p.add_argument("--user", default=DEFAULT_USER, help="Service user")
    p.add_argument("--group", default=DEFAULT_GROUP, help="Service group")
    p.add_argument("--config", default=DEFAULT_CONFIG_PATH, help="Path to Loki config")
    p.add_argument("--data-dir", default=DEFAULT_DATA_DIR, help="Loki data dir")
    p.add_argument("--log-dir", default=DEFAULT_LOG_DIR, help="Loki log dir")
    p.add_argument("--quiet", action="store_true", help="Less output")
    args = p.parse_args()

    quiet = args.quiet or args.update

    if os.geteuid() != 0:
        print("Run as root.", file=sys.stderr)
        return 1

    ensure_user(args.user, args.group, quiet=quiet)

    ensure_dirs([
        os.path.dirname(args.config) or DEFAULT_CONFIG_DIR,
        args.data_dir,
        os.path.join(args.data_dir, "chunks"),
        os.path.join(args.data_dir, "rules"),
        os.path.join(args.data_dir, "db"),
        os.path.join(args.data_dir, "compactor"),
        args.log_dir,
    ], owner=(args.user, args.group), quiet=quiet, force=args.force)

    write_config(args.config, CUSTOM_CONFIG, owner=(args.user, args.group), quiet=quiet, force=args.force)

    installed = get_installed_version(quiet=quiet)
    deb_url = args.deb_url
    if deb_url == "latest":
        try:
            latest_ver, latest_url = find_latest_loki_deb_asset(quiet=quiet)
        except (URLError, HTTPError):
            return 0
        if not latest_ver or not latest_url:
            return 0
        deb_url = latest_url
    if args.update or args.force:
        try:
            latest_ver, latest_url = find_latest_loki_deb_asset(quiet=quiet)
        except (URLError, HTTPError):
            return 0
        if not latest_ver or not latest_url:
            return 0
        if args.force or not installed or semver_tuple(installed) < semver_tuple(latest_ver):
            try:
                install_deb_from_url(latest_url, quiet=True, force=args.force)
            except SystemExit:
                return 1
    else:
        if not installed:
            install_deb_from_url(deb_url, quiet=quiet)

    ensure_systemd(config_path=args.config, user=args.user, group=args.group, quiet=quiet, force=args.force)

    run(f"chown -R {args.user}:{args.group} {os.path.dirname(args.config)} {args.data_dir} {args.log_dir}")

    log("OK.", quiet)
    return 0

if __name__ == "__main__":
    sys.exit(main())