package storage

import (
	"context"
	"fmt"
	"time"

	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"

	"git.tornberg.me/go-otel/internal/honeypot"
	"git.tornberg.me/go-otel/internal/uaclass"
)

// SourceIP is one unique client address ever seen on ingest, with lifetime
// totals. The address is the edge peer when the trace has one (the real
// client behind ingress), otherwise the span's own peer.
type SourceIP struct {
	IP        string    `json:"ip"`
	FirstSeen time.Time `json:"firstSeen"`
	LastSeen  time.Time `json:"lastSeen"`
	Requests  uint64    `json:"requests"`
	Errors    uint64    `json:"errors"`
	Services  []string  `json:"services"`
	UserAgent string    `json:"userAgent,omitempty"`
	UAClass   string    `json:"uaClass,omitempty"`
	// Public is false for loopback/private/reserved addresses, which are
	// usually a proxy hop rather than a real client.
	Public bool `json:"public"`
	// BadActor is the features.MarkBadActors rule the recent requests match
	// ("" if none); ProbePaths counts the distinct probe paths they hit. Set by
	// the caller from SourceIPRoutes, not by SourceIPs.
	BadActor   string `json:"badActor,omitempty"`
	ProbePaths int    `json:"probePaths,omitempty"`
}

// SourceIPEvent is one request attributed to a source address.
type SourceIPEvent struct {
	Timestamp  time.Time `json:"timestamp"`
	Service    string    `json:"service"`
	Method     string    `json:"method,omitempty"`
	Host       string    `json:"host,omitempty"`
	Route      string    `json:"route,omitempty"`
	SpanName   string    `json:"spanName,omitempty"`
	StatusCode uint16    `json:"statusCode"`
	IsError    bool      `json:"isError"`
	TraceID    string    `json:"traceId"`
	SpanID     string    `json:"spanId"`
	UserAgent  string    `json:"userAgent,omitempty"`
}

// SourceIPQuerier is implemented by stores that keep the global source
// address list (ClickHouseStore).
type SourceIPQuerier interface {
	// SourceIPs lists unique addresses, most recently seen first. search is
	// a substring match on the address; since (zero = all time) drops
	// addresses last seen earlier.
	SourceIPs(ctx context.Context, search string, since time.Time, limit int) ([]SourceIP, error)
	// SourceIPActivity lists the requests attributed to one address, newest first.
	SourceIPActivity(ctx context.Context, ip string, since time.Time, limit int) ([]SourceIPEvent, error)
	// SourceIPRoutes returns, per address, the distinct (route, status, user
	// agent) combinations requested since the given time - enough to judge
	// scanning behaviour for many addresses in one query.
	SourceIPRoutes(ctx context.Context, ips []string, since time.Time) (map[string][]SourceIPEvent, error)
}

// sourceIPClient is the client-address expression shared by the rollup and
// the activity query, so both agree on who a request "came from".
const sourceIPClient = "if(edge_peer_ip != '', edge_peer_ip, peer_ip)"

// sourceIPRequest keeps one row per request: a request produces several
// spans (ingress, upstream, ...) that all carry the same edge peer, so only
// the edge span itself is counted.
const sourceIPRequest = "kind = 'span' AND ip != '' AND (edge_peer_ip = '' OR span_id = edge_span_id)"

const sourceIPSelect = `SELECT %s AS ip, min(timestamp) AS first_seen, max(timestamp) AS last_seen,
	count() AS requests, sum(is_error) AS errors, groupUniqArray(service) AS services,
	anyLast(user_agent) AS user_agent, anyLast(ua_class) AS ua_class
	FROM %s.normalized_events WHERE %s GROUP BY ip`

// EnsureSourceIPSchema creates the lifetime source_ips table and the
// materialized view feeding it from normalized_events. Unlike
// normalized_events it has no TTL: the list is meant to outlive raw data.
// An empty table is backfilled once from whatever normalized_events holds.
// Idempotent.
func EnsureSourceIPSchema(ctx context.Context, conn driver.Conn, db string) error {
	if err := conn.Exec(ctx, fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s.source_ips (
		ip String,
		first_seen SimpleAggregateFunction(min, DateTime64(9, 'UTC')),
		last_seen SimpleAggregateFunction(max, DateTime64(9, 'UTC')),
		requests SimpleAggregateFunction(sum, UInt64),
		errors SimpleAggregateFunction(sum, UInt64),
		services SimpleAggregateFunction(groupUniqArrayArray, Array(String)),
		user_agent SimpleAggregateFunction(anyLast, String),
		ua_class SimpleAggregateFunction(anyLast, String)
	) ENGINE = AggregatingMergeTree
	ORDER BY ip`, db)); err != nil {
		return err
	}
	sel := fmt.Sprintf(sourceIPSelect, sourceIPClient, db, sourceIPRequest)

	var existing uint64
	if err := conn.QueryRow(ctx, fmt.Sprintf("SELECT count() FROM %s.source_ips", db)).Scan(&existing); err != nil {
		return err
	}
	if existing == 0 {
		if err := conn.Exec(ctx, fmt.Sprintf("INSERT INTO %s.source_ips %s", db, sel)); err != nil {
			return fmt.Errorf("backfill: %w", err)
		}
	}
	return conn.Exec(ctx, fmt.Sprintf("CREATE MATERIALIZED VIEW IF NOT EXISTS %s.source_ips_mv TO %s.source_ips AS %s", db, db, sel))
}

func (s *ClickHouseStore) SourceIPs(ctx context.Context, search string, since time.Time, limit int) ([]SourceIP, error) {
	if limit <= 0 || limit > 1000 {
		limit = 200
	}
	where, args := "1", []interface{}{}
	if search != "" {
		where += " AND position(ip, ?) > 0"
		args = append(args, search)
	}
	having := ""
	if !since.IsZero() {
		having = "HAVING last_seen >= ?"
		args = append(args, since.UTC())
	}
	args = append(args, limit)
	rows, err := s.conn.Query(ctx, fmt.Sprintf(`SELECT ip, min(first_seen), max(last_seen) AS last_seen, sum(requests), sum(errors),
		groupUniqArrayArray(services), anyLast(user_agent), anyLast(ua_class)
		FROM %s.source_ips WHERE %s GROUP BY ip %s ORDER BY last_seen DESC LIMIT ?`, s.db, where, having), args...)
	if err != nil {
		return nil, fmt.Errorf("clickhouse source ips: %w", err)
	}
	defer rows.Close()
	out := []SourceIP{}
	for rows.Next() {
		var v SourceIP
		if err := rows.Scan(&v.IP, &v.FirstSeen, &v.LastSeen, &v.Requests, &v.Errors, &v.Services, &v.UserAgent, &v.UAClass); err != nil {
			return nil, err
		}
		if v.UAClass == "" && v.UserAgent != "" {
			v.UAClass = uaclass.Of(v.UserAgent)
		}
		v.Public = !honeypot.Untrackable(v.IP)
		out = append(out, v)
	}
	return out, rows.Err()
}

func (s *ClickHouseStore) SourceIPActivity(ctx context.Context, ip string, since time.Time, limit int) ([]SourceIPEvent, error) {
	if limit <= 0 || limit > 1000 {
		limit = 200
	}
	if since.IsZero() {
		since = time.Now().Add(-7 * 24 * time.Hour)
	}
	rows, err := s.conn.Query(ctx, fmt.Sprintf(`SELECT timestamp, service, method, host, route, span_name, status_code, is_error, trace_id, span_id, user_agent
		FROM (SELECT *, %s AS ip FROM %s.normalized_events WHERE timestamp >= ?)
		WHERE %s AND ip = ? ORDER BY timestamp DESC LIMIT ?`, sourceIPClient, s.db, sourceIPRequest), since.UTC(), ip, limit)
	if err != nil {
		return nil, fmt.Errorf("clickhouse source ip activity: %w", err)
	}
	defer rows.Close()
	out := []SourceIPEvent{}
	for rows.Next() {
		var v SourceIPEvent
		var isErr uint8
		if err := rows.Scan(&v.Timestamp, &v.Service, &v.Method, &v.Host, &v.Route, &v.SpanName, &v.StatusCode, &isErr, &v.TraceID, &v.SpanID, &v.UserAgent); err != nil {
			return nil, err
		}
		v.IsError = isErr != 0
		out = append(out, v)
	}
	return out, rows.Err()
}

func (s *ClickHouseStore) SourceIPRoutes(ctx context.Context, ips []string, since time.Time) (map[string][]SourceIPEvent, error) {
	out := map[string][]SourceIPEvent{}
	if len(ips) == 0 {
		return out, nil
	}
	rows, err := s.conn.Query(ctx, fmt.Sprintf(`SELECT ip, route, status_code, anyLast(user_agent)
		FROM (SELECT *, %s AS ip FROM %s.normalized_events WHERE timestamp >= ? AND kind = 'span')
		WHERE %s AND ip IN ? GROUP BY ip, route, status_code`, sourceIPClient, s.db, sourceIPRequest), since.UTC(), ips)
	if err != nil {
		return nil, fmt.Errorf("clickhouse source ip routes: %w", err)
	}
	defer rows.Close()
	for rows.Next() {
		var ip string
		var v SourceIPEvent
		if err := rows.Scan(&ip, &v.Route, &v.StatusCode, &v.UserAgent); err != nil {
			return nil, err
		}
		out[ip] = append(out[ip], v)
	}
	return out, rows.Err()
}