package storage

import (
	"context"
	"encoding/json"
	"fmt"
	"sort"
	"time"

	"git.tornberg.me/go-otel/internal/websocket"
)

// initTraceSchema adds everything the trace and issue overviews need on top of
// the raw entries table:
//
//   - entries.trace_ids + bloom indexes: indexed "which payloads contain trace X"
//     instead of a substring scan over every payload.
//   - span_facts (Null engine) -> trace_summaries (AggregatingMergeTree): one
//     row per trace maintained on insert, so overview tables never read spans.
//   - issue_occurrences: every occurrence of every issue, beyond the 5 samples
//     the tracker keeps in memory.
//
// Failures are logged, not fatal: raw storage keeps working and the API falls
// back to the slower paths.
func (s *ClickHouseStore) initTraceSchema(ctx context.Context) error {
	ddl := []string{
		fmt.Sprintf(`ALTER TABLE %s.entries ADD COLUMN IF NOT EXISTS trace_ids Array(String) DEFAULT []`, s.db),

		fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s.span_facts (
			at             DateTime('UTC'),
			trace_id       String,
			span_id        String,
			parent_span_id String,
			service        LowCardinality(String),
			name           String,
			t0_ns          UInt64,
			t1_ns          UInt64,
			is_error       UInt8
		) ENGINE = Null`, s.db),

		fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s.trace_summaries (
			trace_id String,
			ts       SimpleAggregateFunction(min, DateTime('UTC')),
			start_ns SimpleAggregateFunction(min, UInt64),
			end_ns   SimpleAggregateFunction(max, UInt64),
			spans    SimpleAggregateFunction(sum, UInt64),
			errors   SimpleAggregateFunction(sum, UInt64),
			services SimpleAggregateFunction(groupUniqArrayArray, Array(String)),
			root     AggregateFunction(argMin, String, Tuple(UInt8, UInt64))
		) ENGINE = AggregatingMergeTree
		ORDER BY trace_id
		TTL ts + INTERVAL 90 DAY`, s.db),

		fmt.Sprintf(`CREATE MATERIALIZED VIEW IF NOT EXISTS %s.trace_summaries_mv TO %s.trace_summaries AS
			SELECT
				trace_id,
				min(at)                                              AS ts,
				min(t0_ns)                                           AS start_ns,
				max(t1_ns)                                           AS end_ns,
				count()                                              AS spans,
				sum(is_error)                                        AS errors,
				groupUniqArray(service)                              AS services,
				argMinState(name, (parent_span_id != '', t0_ns))  AS root
			FROM %s.span_facts
			GROUP BY trace_id`, s.db, s.db, s.db),

		fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s.issue_occurrences (
			issue_id  String,
			timestamp DateTime64(9, 'UTC'),
			kind      LowCardinality(String),
			service   LowCardinality(String),
			trace_id  String,
			span_id   String,
			sample    String CODEC(ZSTD(3))
		) ENGINE = MergeTree
		PARTITION BY toYYYYMM(timestamp)
		ORDER BY (issue_id, timestamp)
		TTL toDateTime(timestamp) + INTERVAL 90 DAY`, s.db),
	}
	for _, q := range ddl {
		if err := s.conn.Exec(ctx, q); err != nil {
			fmt.Printf("clickhouse: trace schema step failed (falling back to slow paths): %v\n", err)
			return nil
		}
	}

	for _, col := range []string{"trace_id", "trace_ids"} {
		s.ensureBloomIndex(ctx, "entries", "idx_"+col, col)
	}
	return nil
}

// ensureBloomIndex adds a bloom-filter skip index and, only when it is newly
// added, materializes it over existing parts (a background mutation).
func (s *ClickHouseStore) ensureBloomIndex(ctx context.Context, table, name, expr string) {
	var n uint64
	row := s.conn.QueryRow(ctx,
		`SELECT count() FROM system.data_skipping_indices WHERE database = ? AND table = ? AND name = ?`,
		s.db, table, name)
	if err := row.Scan(&n); err != nil || n > 0 {
		return
	}
	if err := s.conn.Exec(ctx, fmt.Sprintf(
		`ALTER TABLE %s.%s ADD INDEX IF NOT EXISTS %s %s TYPE bloom_filter(0.01) GRANULARITY 1`,
		s.db, table, name, expr)); err != nil {
		fmt.Printf("clickhouse: add index %s: %v\n", name, err)
		return
	}
	if err := s.conn.Exec(ctx, fmt.Sprintf(`ALTER TABLE %s.%s MATERIALIZE INDEX %s`, s.db, table, name)); err != nil {
		fmt.Printf("clickhouse: materialize index %s: %v\n", name, err)
	}
}

func (s *ClickHouseStore) appendSpanFacts(ctx context.Context, ts time.Time, facts []spanFact) error {
	if len(facts) == 0 {
		return nil
	}
	batch, err := s.conn.PrepareBatch(ctx, fmt.Sprintf(
		`INSERT INTO %s.span_facts (at, trace_id, span_id, parent_span_id, service, name, t0_ns, t1_ns, is_error)`, s.db))
	if err != nil {
		return err
	}
	for _, f := range facts {
		if err := batch.Append(ts, f.TraceID, f.SpanID, f.ParentSpanID, f.Service, f.Name, f.StartNs, f.EndNs, boolToUint8(f.IsError)); err != nil {
			_ = batch.Abort()
			return err
		}
	}
	return batch.Send()
}

// TraceEntries implements TraceQuerier.
func (s *ClickHouseStore) TraceEntries(ctx context.Context, traceID string, limit int) ([]websocket.HistoryEntry, error) {
	if limit <= 0 {
		limit = 1000
	}
	rows, err := s.conn.Query(ctx, fmt.Sprintf(
		`SELECT timestamp, type, data FROM %s.entries
		 WHERE has(trace_ids, ?) OR trace_id = ?
		 ORDER BY timestamp DESC LIMIT ?`, s.db), traceID, traceID, limit)
	if err != nil {
		return nil, fmt.Errorf("clickhouse trace entries: %w", err)
	}
	defer rows.Close()

	var out []websocket.HistoryEntry
	for rows.Next() {
		var (
			ts  time.Time
			typ string
			raw string
		)
		if err := rows.Scan(&ts, &typ, &raw); err != nil {
			return nil, err
		}
		trimmed, ok := TrimToTrace([]byte(raw), traceID)
		if !ok {
			continue
		}
		out = append(out, websocket.HistoryEntry{Type: typ, Timestamp: ts.UnixNano(), Data: json.RawMessage(trimmed)})
	}
	return out, rows.Err()
}

// TraceSummaries implements TraceQuerier.
func (s *ClickHouseStore) TraceSummaries(ctx context.Context, ids []string) (map[string]TraceSummary, error) {
	out := make(map[string]TraceSummary, len(ids))
	if len(ids) == 0 {
		return out, nil
	}
	rows, err := s.conn.Query(ctx, fmt.Sprintf(
		`SELECT trace_id, min(start_ns), max(end_ns), sum(spans), sum(errors),
		        groupUniqArrayArray(services), argMinMerge(root)
		 FROM %s.trace_summaries
		 WHERE trace_id IN ?
		 GROUP BY trace_id`, s.db), ids)
	if err != nil {
		return nil, fmt.Errorf("clickhouse trace summaries: %w", err)
	}
	defer rows.Close()
	for rows.Next() {
		var (
			id, name             string
			start, end, spans, e uint64
			services             []string
		)
		if err := rows.Scan(&id, &start, &end, &spans, &e, &services, &name); err != nil {
			return nil, err
		}
		sort.Strings(services)
		if name == "" {
			name = "Trace " + id[:min(8, len(id))]
		}
		out[id] = TraceSummary{
			ID: id, Name: name, StartTime: int64(start), EndTime: int64(end), Duration: int64(end - start),
			SpanCount: int(spans), ErrorCount: int(e), Services: services,
		}
	}
	return out, rows.Err()
}

// AppendOccurrence implements OccurrenceStore.
func (s *ClickHouseStore) AppendOccurrence(ctx context.Context, occ Occurrence) error {
	err := s.conn.AsyncInsert(ctx, fmt.Sprintf(
		`INSERT INTO %s.issue_occurrences (issue_id, timestamp, kind, service, trace_id, span_id, sample)
		 VALUES (?, ?, ?, ?, ?, ?, ?)`, s.db),
		false,
		occ.IssueID, time.Unix(0, occ.Timestamp).UTC(), occ.Kind, occ.Service, occ.TraceID, occ.SpanID, string(occ.Sample),
	)
	if err != nil {
		return fmt.Errorf("clickhouse append occurrence: %w", err)
	}
	return nil
}

const maxSamplesPerIssueTrace = 20
const maxUntraced = 20

// IssueTraces implements OccurrenceStore: one aggregate query for the page of
// traces, one for their samples, one batched lookup for their summaries.
func (s *ClickHouseStore) IssueTraces(ctx context.Context, issueID string, opts IssueTracesOptions) (IssueTracesPage, error) {
	limit := opts.Limit
	if limit <= 0 || limit > 200 {
		limit = 50
	}
	var page IssueTracesPage

	var total uint64
	if err := s.conn.QueryRow(ctx, fmt.Sprintf(
		`SELECT count() FROM %s.issue_occurrences WHERE issue_id = ?`, s.db), issueID).Scan(&total); err != nil {
		return page, fmt.Errorf("clickhouse issue totals: %w", err)
	}
	page.TotalOccurrences = int(total)

	having := ""
	args := []any{issueID}
	if opts.Before > 0 {
		having = "HAVING max(timestamp) < ?"
		args = append(args, time.Unix(0, opts.Before).UTC())
	}
	args = append(args, limit+1)
	rows, err := s.conn.Query(ctx, fmt.Sprintf(
		`SELECT trace_id, count(), min(timestamp), max(timestamp)
		 FROM %s.issue_occurrences
		 WHERE issue_id = ? AND trace_id != ''
		 GROUP BY trace_id %s
		 ORDER BY max(timestamp) DESC
		 LIMIT ?`, s.db, having), args...)
	if err != nil {
		return page, fmt.Errorf("clickhouse issue traces: %w", err)
	}
	for rows.Next() {
		var (
			t        IssueTrace
			hits     uint64
			from, to time.Time
		)
		if err := rows.Scan(&t.TraceID, &hits, &from, &to); err != nil {
			_ = rows.Close()
			return page, err
		}
		t.Hits, t.FirstSeen, t.LastSeen = int(hits), from.UnixNano(), to.UnixNano()
		page.Traces = append(page.Traces, t)
	}
	_ = rows.Close()
	if err := rows.Err(); err != nil {
		return page, err
	}
	if len(page.Traces) > limit {
		page.Traces = page.Traces[:limit]
		page.HasMore = true
	}

	if len(page.Traces) > 0 {
		ids := make([]string, len(page.Traces))
		index := make(map[string]int, len(ids))
		for i, t := range page.Traces {
			ids[i] = t.TraceID
			index[t.TraceID] = i
		}
		srows, err := s.conn.Query(ctx, fmt.Sprintf(
			`SELECT trace_id, sample FROM %s.issue_occurrences
			 WHERE issue_id = ? AND trace_id IN ?
			 ORDER BY timestamp DESC
			 LIMIT ? BY trace_id`, s.db), issueID, ids, maxSamplesPerIssueTrace)
		if err != nil {
			return page, fmt.Errorf("clickhouse issue trace samples: %w", err)
		}
		for srows.Next() {
			var id, sample string
			if err := srows.Scan(&id, &sample); err != nil {
				_ = srows.Close()
				return page, err
			}
			i := index[id]
			page.Traces[i].Samples = append(page.Traces[i].Samples, json.RawMessage(sample))
		}
		_ = srows.Close()
		if err := srows.Err(); err != nil {
			return page, err
		}

		summaries, err := s.TraceSummaries(ctx, ids)
		if err != nil {
			return page, err
		}
		for i := range page.Traces {
			if sum, ok := summaries[page.Traces[i].TraceID]; ok {
				page.Traces[i].Summary = &sum
			}
		}
	}

	if opts.Before == 0 {
		urows, err := s.conn.Query(ctx, fmt.Sprintf(
			`SELECT sample FROM %s.issue_occurrences
			 WHERE issue_id = ? AND trace_id = ''
			 ORDER BY timestamp DESC LIMIT ?`, s.db), issueID, maxUntraced)
		if err != nil {
			return page, fmt.Errorf("clickhouse issue untraced: %w", err)
		}
		for urows.Next() {
			var sample string
			if err := urows.Scan(&sample); err != nil {
				_ = urows.Close()
				return page, err
			}
			page.Untraced = append(page.Untraced, json.RawMessage(sample))
		}
		_ = urows.Close()
		if err := urows.Err(); err != nil {
			return page, err
		}
	}
	return page, nil
}

// IssueSources returns every distinct source IP recorded for an issue. Source
// addresses may be in the derived occurrence sample, or on any normalized span
// belonging to one of the issue's traces (for example, an ingress span's
// net.peer.ip when the failing child span has no client attributes).
func (s *ClickHouseStore) IssueSources(ctx context.Context, issueID string) ([]IssueSource, error) {
	rows, err := s.conn.Query(ctx, fmt.Sprintf(`
		SELECT ip, min(seen_at), max(seen_at), count()
		FROM (
			-- Logs and older occurrences keep their source address in the
			-- derived sample. Keep this path for data that has no trace.
			SELECT JSONExtractString(sample, 'sourceIp') AS ip, timestamp AS seen_at
			FROM %s.issue_occurrences
			WHERE issue_id = ? AND ip != ''

			UNION ALL

			-- A span error's sample describes the failing span, but the public
			-- peer is commonly recorded on the ingress span in the same trace.
			-- Include every normalized span peer for the issue's traces.
			SELECT peer_ip AS ip, timestamp AS seen_at
			FROM %s.normalized_events
			WHERE peer_ip != ''
			  AND trace_id IN (
				SELECT DISTINCT trace_id
				FROM %s.issue_occurrences
				WHERE issue_id = ? AND trace_id != ''
			  )
		)
		WHERE ip != ''
		GROUP BY ip
		ORDER BY count() DESC, max(seen_at) DESC`, s.db, s.db, s.db), issueID, issueID)
	if err != nil {
		return nil, fmt.Errorf("clickhouse issue sources: %w", err)
	}
	defer rows.Close()
	var out []IssueSource
	for rows.Next() {
		var source IssueSource
		if err := rows.Scan(&source.IP, &source.FirstSeen, &source.LastSeen, &source.Count); err != nil {
			return nil, err
		}
		out = append(out, source)
	}
	return out, rows.Err()
}