package metrics

import (
	"context"
	"net/http"

	"code.t-juice.club/torjus/oubliette/internal/storage"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/collectors"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

// Metrics holds all Prometheus collectors for the honeypot.
type Metrics struct {
	registry *prometheus.Registry

	SSHConnectionsTotal    *prometheus.CounterVec
	SSHConnectionsActive   prometheus.Gauge
	AuthAttemptsTotal      *prometheus.CounterVec
	AuthAttemptsByCountry  *prometheus.CounterVec
	CommandsExecuted       *prometheus.CounterVec
	HumanScore             prometheus.Histogram
	SessionsTotal          *prometheus.CounterVec
	SessionsActive         prometheus.Gauge
	SessionDuration        prometheus.Histogram
	ExecCommandsTotal      prometheus.Counter
	BuildInfo              *prometheus.GaugeVec
	StorageQueryDuration   *prometheus.HistogramVec
	StorageQueryErrors     *prometheus.CounterVec
}

// New creates a new Metrics instance with all collectors registered.
func New(version string) *Metrics {
	reg := prometheus.NewRegistry()

	m := &Metrics{
		registry: reg,
		SSHConnectionsTotal: prometheus.NewCounterVec(prometheus.CounterOpts{
			Name: "oubliette_ssh_connections_total",
			Help: "Total SSH connections received.",
		}, []string{"outcome"}),
		SSHConnectionsActive: prometheus.NewGauge(prometheus.GaugeOpts{
			Name: "oubliette_ssh_connections_active",
			Help: "Current active SSH connections.",
		}),
		AuthAttemptsTotal: prometheus.NewCounterVec(prometheus.CounterOpts{
			Name: "oubliette_auth_attempts_total",
			Help: "Total authentication attempts.",
		}, []string{"result", "reason"}),
		AuthAttemptsByCountry: prometheus.NewCounterVec(prometheus.CounterOpts{
			Name: "oubliette_auth_attempts_by_country_total",
			Help: "Total authentication attempts by country.",
		}, []string{"country"}),
		CommandsExecuted: prometheus.NewCounterVec(prometheus.CounterOpts{
			Name: "oubliette_commands_executed_total",
			Help: "Total commands executed in shells.",
		}, []string{"shell"}),
		HumanScore: prometheus.NewHistogram(prometheus.HistogramOpts{
			Name:    "oubliette_human_score",
			Help:    "Distribution of final human detection scores.",
			Buckets: prometheus.LinearBuckets(0, 0.1, 11), // 0.0, 0.1, ..., 1.0
		}),
		SessionsTotal: prometheus.NewCounterVec(prometheus.CounterOpts{
			Name: "oubliette_sessions_total",
			Help: "Total sessions created.",
		}, []string{"shell"}),
		SessionsActive: prometheus.NewGauge(prometheus.GaugeOpts{
			Name: "oubliette_sessions_active",
			Help: "Current active sessions.",
		}),
		SessionDuration: prometheus.NewHistogram(prometheus.HistogramOpts{
			Name:    "oubliette_session_duration_seconds",
			Help:    "Session duration in seconds.",
			Buckets: []float64{1, 5, 10, 30, 60, 120, 300, 600, 1800, 3600},
		}),
		ExecCommandsTotal: prometheus.NewCounter(prometheus.CounterOpts{
			Name: "oubliette_exec_commands_total",
			Help: "Total SSH exec commands received.",
		}),
		BuildInfo: prometheus.NewGaugeVec(prometheus.GaugeOpts{
			Name: "oubliette_build_info",
			Help: "Build information. Always 1.",
		}, []string{"version"}),
		StorageQueryDuration: prometheus.NewHistogramVec(prometheus.HistogramOpts{
			Name:    "oubliette_storage_query_duration_seconds",
			Help:    "Duration of storage query calls in seconds.",
			Buckets: []float64{0.001, 0.005, 0.01, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10},
		}, []string{"method"}),
		StorageQueryErrors: prometheus.NewCounterVec(prometheus.CounterOpts{
			Name: "oubliette_storage_query_errors_total",
			Help: "Total storage query errors.",
		}, []string{"method"}),
	}

	reg.MustRegister(
		collectors.NewGoCollector(),
		collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}),
		m.SSHConnectionsTotal,
		m.SSHConnectionsActive,
		m.AuthAttemptsTotal,
		m.AuthAttemptsByCountry,
		m.CommandsExecuted,
		m.HumanScore,
		m.SessionsTotal,
		m.SessionsActive,
		m.SessionDuration,
		m.ExecCommandsTotal,
		m.BuildInfo,
		m.StorageQueryDuration,
		m.StorageQueryErrors,
	)

	m.BuildInfo.WithLabelValues(version).Set(1)

	// Initialize label combinations so they appear in Gather/output.
	for _, outcome := range []string{"accepted", "rejected_handshake", "rejected_max_connections"} {
		m.SSHConnectionsTotal.WithLabelValues(outcome)
	}
	for _, reason := range []string{"static_credential", "remembered_credential", "threshold_reached", "rejected"} {
		m.AuthAttemptsTotal.WithLabelValues("accepted", reason)
		m.AuthAttemptsTotal.WithLabelValues("rejected", reason)
	}
	for _, sh := range []string{"bash", "fridge", "banking", "adventure", "cisco"} {
		m.SessionsTotal.WithLabelValues(sh)
		m.CommandsExecuted.WithLabelValues(sh)
	}

	return m
}

// RegisterStoreCollector registers a collector that queries storage stats on each scrape.
func (m *Metrics) RegisterStoreCollector(store storage.Store) {
	m.registry.MustRegister(&storeCollector{store: store})
}

// Handler returns an http.Handler that serves Prometheus metrics.
func (m *Metrics) Handler() http.Handler {
	return promhttp.HandlerFor(m.registry, promhttp.HandlerOpts{})
}

// storeCollector implements prometheus.Collector, querying storage on each scrape.
type storeCollector struct {
	store storage.Store
}

var (
	storageLoginAttemptsDesc = prometheus.NewDesc(
		"oubliette_storage_login_attempts_total",
		"Total login attempts in storage.",
		nil, nil,
	)
	storageUniqueIPsDesc = prometheus.NewDesc(
		"oubliette_storage_unique_ips",
		"Unique IPs in storage.",
		nil, nil,
	)
	storageSessionsDesc = prometheus.NewDesc(
		"oubliette_storage_sessions_total",
		"Total sessions in storage.",
		nil, nil,
	)
)

func (c *storeCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- storageLoginAttemptsDesc
	ch <- storageUniqueIPsDesc
	ch <- storageSessionsDesc
}

func (c *storeCollector) Collect(ch chan<- prometheus.Metric) {
	stats, err := c.store.GetDashboardStats(context.Background())
	if err != nil {
		return
	}
	ch <- prometheus.MustNewConstMetric(storageLoginAttemptsDesc, prometheus.GaugeValue, float64(stats.TotalAttempts))
	ch <- prometheus.MustNewConstMetric(storageUniqueIPsDesc, prometheus.GaugeValue, float64(stats.UniqueIPs))
	ch <- prometheus.MustNewConstMetric(storageSessionsDesc, prometheus.GaugeValue, float64(stats.TotalSessions))
}