{ config, lib, pkgs, ... }:

let
  cfg = config.services.burrow.observability;
  otelCollectorPackage =
    if builtins.hasAttr "opentelemetry-collector-contrib" pkgs then
      pkgs."opentelemetry-collector-contrib"
    else
      pkgs.opentelemetry-collector;
  garageCfg = config.services.burrow.garage;
  headscaleCfg = config.services.burrow.headscale;
  grafanaHttpPort = config.services.grafana.settings.server.http_port or 3000;
  grafanaHttpAddress = config.services.grafana.settings.server.http_addr or "127.0.0.1";
  tailscaleExporterEnabled = cfg.tailscaleExporterEnvironmentFile != null;
in
{
  options.services.burrow.observability = {
    enable = lib.mkEnableOption "Burrow Prometheus, OpenTelemetry, Jaeger, and Grafana wiring";

    prometheusListenAddress = lib.mkOption {
      type = lib.types.str;
      default = "127.0.0.1";
      description = "Address used by the local Prometheus server.";
    };

    prometheusPort = lib.mkOption {
      type = lib.types.port;
      default = 9090;
      description = "Port used by the local Prometheus server.";
    };

    otelGrpcEndpoint = lib.mkOption {
      type = lib.types.str;
      default = "127.0.0.1:4317";
      description = "Local OTLP gRPC endpoint accepted by the OpenTelemetry collector.";
    };

    otelHttpEndpoint = lib.mkOption {
      type = lib.types.str;
      default = "127.0.0.1:4318";
      description = "Local OTLP HTTP endpoint accepted by the OpenTelemetry collector.";
    };

    otelPrometheusEndpoint = lib.mkOption {
      type = lib.types.str;
      default = "127.0.0.1:9464";
      description = "Prometheus endpoint where the OpenTelemetry collector exports its own metrics.";
    };

    jaegerImage = lib.mkOption {
      type = lib.types.str;
      default = "docker.io/jaegertracing/all-in-one:1.57";
      description = "OCI image used for the local Jaeger all-in-one service.";
    };

    jaegerUiPort = lib.mkOption {
      type = lib.types.port;
      default = 16686;
      description = "Local Jaeger UI and query API port.";
    };

    jaegerOtlpGrpcPort = lib.mkOption {
      type = lib.types.port;
      default = 14317;
      description = "Local remapped OTLP gRPC port forwarded to Jaeger.";
    };

    jaegerOtlpHttpPort = lib.mkOption {
      type = lib.types.port;
      default = 14318;
      description = "Local remapped OTLP HTTP port forwarded to Jaeger.";
    };

    jaegerAdminPort = lib.mkOption {
      type = lib.types.port;
      default = 14269;
      description = "Local Jaeger admin and metrics port.";
    };

    tailscaleExporterEnvironmentFile = lib.mkOption {
      type = lib.types.nullOr lib.types.path;
      default = null;
      description = "Optional environment file for the Prometheus Tailscale exporter OAuth credentials.";
    };
  };

  config = lib.mkIf cfg.enable (lib.mkMerge [
    {
      virtualisation.podman.enable = true;

      services.grafana = {
        settings = {
          metrics = {
            enabled = true;
            disable_total_stats = true;
          };
        };
        provision = {
          enable = true;
          datasources.settings = {
            apiVersion = 1;
            prune = true;
            datasources = [
              {
                name = "Prometheus";
                uid = "prometheus";
                type = "prometheus";
                access = "proxy";
                url = "http://${cfg.prometheusListenAddress}:${toString cfg.prometheusPort}";
                isDefault = true;
                jsonData = {
                  timeInterval = "15s";
                  httpMethod = "POST";
                };
              }
              {
                name = "Jaeger";
                uid = "jaeger";
                type = "jaeger";
                access = "proxy";
                url = "http://127.0.0.1:${toString cfg.jaegerUiPort}";
              }
            ];
          };
        };
      };

      services.prometheus = {
        enable = true;
        listenAddress = cfg.prometheusListenAddress;
        port = cfg.prometheusPort;
        retentionTime = "30d";
        checkConfig = "syntax-only";
        exporters = {
          node = {
            enable = true;
            listenAddress = "127.0.0.1";
          };
          systemd = {
            enable = true;
            listenAddress = "127.0.0.1";
          };
        };
        scrapeConfigs = [
          {
            job_name = "prometheus";
            static_configs = [
              { targets = [ "${cfg.prometheusListenAddress}:${toString cfg.prometheusPort}" ]; }
            ];
          }
          {
            job_name = "node";
            static_configs = [
              { targets = [ "127.0.0.1:${toString config.services.prometheus.exporters.node.port}" ]; }
            ];
          }
          {
            job_name = "systemd";
            static_configs = [
              { targets = [ "127.0.0.1:${toString config.services.prometheus.exporters.systemd.port}" ]; }
            ];
          }
          {
            job_name = "grafana";
            metrics_path = "/metrics";
            static_configs = [
              { targets = [ "${grafanaHttpAddress}:${toString grafanaHttpPort}" ]; }
            ];
          }
          {
            job_name = "headscale";
            metrics_path = "/metrics";
            static_configs = [
              { targets = [ "${headscaleCfg.metricsListenAddress}:${toString headscaleCfg.metricsPort}" ]; }
            ];
          }
          {
            job_name = "otel-collector";
            static_configs = [
              { targets = [ cfg.otelPrometheusEndpoint ]; }
            ];
          }
          {
            job_name = "jaeger";
            metrics_path = "/metrics";
            static_configs = [
              { targets = [ "127.0.0.1:${toString cfg.jaegerAdminPort}" ]; }
            ];
          }
        ] ++ lib.optionals garageCfg.enable [
          {
            job_name = "garage";
            metrics_path = "/metrics";
            static_configs = [
              { targets = [ "${garageCfg.adminListenAddress}:${toString garageCfg.adminPort}" ]; }
            ];
          }
        ] ++ lib.optionals tailscaleExporterEnabled [
          {
            job_name = "tailscale";
            static_configs = [
              { targets = [ "127.0.0.1:${toString config.services.prometheus.exporters.tailscale.port}" ]; }
            ];
          }
        ];
      };

      services.opentelemetry-collector = {
        enable = true;
        package = otelCollectorPackage;
        settings = {
          receivers = {
            otlp = {
              protocols = {
                grpc.endpoint = cfg.otelGrpcEndpoint;
                http.endpoint = cfg.otelHttpEndpoint;
              };
            };
          };
          processors = {
            resource.attributes = [
              {
                key = "service.namespace";
                value = "burrow";
                action = "upsert";
              }
              {
                key = "deployment.environment";
                value = config.networking.hostName;
                action = "upsert";
              }
            ];
            batch = { };
          };
          exporters = {
            "otlp/jaeger" = {
              endpoint = "127.0.0.1:${toString cfg.jaegerOtlpGrpcPort}";
              tls.insecure = true;
            };
            prometheus.endpoint = cfg.otelPrometheusEndpoint;
          };
          service.pipelines = {
            traces = {
              receivers = [ "otlp" ];
              processors = [
                "resource"
                "batch"
              ];
              exporters = [ "otlp/jaeger" ];
            };
            metrics = {
              receivers = [ "otlp" ];
              processors = [
                "resource"
                "batch"
              ];
              exporters = [ "prometheus" ];
            };
          };
        };
      };

      virtualisation.oci-containers.containers.burrow-jaeger = {
        image = cfg.jaegerImage;
        autoStart = true;
        environment = {
          COLLECTOR_OTLP_ENABLED = "true";
          SPAN_STORAGE_TYPE = "badger";
          BADGER_EPHEMERAL = "false";
          BADGER_DIRECTORY_VALUE = "/badger/data";
          BADGER_DIRECTORY_KEY = "/badger/key";
        };
        volumes = [
          "burrow-jaeger-badger:/badger"
        ];
        ports = [
          "127.0.0.1:${toString cfg.jaegerUiPort}:16686"
          "127.0.0.1:${toString cfg.jaegerOtlpGrpcPort}:4317"
          "127.0.0.1:${toString cfg.jaegerOtlpHttpPort}:4318"
          "127.0.0.1:${toString cfg.jaegerAdminPort}:14269"
        ];
        extraOptions = [ "--pull=always" ];
      };

      systemd.services.opentelemetry-collector.after = [ "podman-burrow-jaeger.service" ];
      systemd.services.opentelemetry-collector.wants = [ "podman-burrow-jaeger.service" ];

      environment.sessionVariables = {
        OTEL_EXPORTER_OTLP_ENDPOINT = "http://${cfg.otelGrpcEndpoint}";
        OTEL_EXPORTER_OTLP_PROTOCOL = "grpc";
      };
    }

    (lib.mkIf tailscaleExporterEnabled {
      services.prometheus.exporters.tailscale = {
        enable = true;
        listenAddress = "127.0.0.1";
        environmentFile = cfg.tailscaleExporterEnvironmentFile;
      };
    })
  ]);
}