// training-sample exports a bounded, time-windowed JSONL dataset from
// normalized_events. It never writes to ClickHouse.
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/ClickHouse/clickhouse-go/v2"
)

type sample struct {
	EventID          string    `json:"eventId"`
	Timestamp        time.Time `json:"timestamp"`
	Service          string    `json:"service"`
	Release          string    `json:"release,omitempty"`
	Kind             string    `json:"kind"`
	Source           string    `json:"source"`
	IsError          uint8     `json:"isError"`
	StatusClass      string    `json:"statusClass,omitempty"`
	StatusCode       uint16    `json:"statusCode,omitempty"`
	DurationMS       float64   `json:"durationMs,omitempty"`
	Route            string    `json:"route,omitempty"`
	SpanName         string    `json:"spanName,omitempty"`
	MetricName       string    `json:"metricName,omitempty"`
	MetricType       string    `json:"metricType,omitempty"`
	MetricValue      float64   `json:"metricValue,omitempty"`
	MetricCount      uint64    `json:"metricCount,omitempty"`
	MetricSum        float64   `json:"metricSum,omitempty"`
	MetricMin        float64   `json:"metricMin,omitempty"`
	MetricMax        float64   `json:"metricMax,omitempty"`
	ExceptionType    string    `json:"exceptionType,omitempty"`
	Message          string    `json:"messageNormalized,omitempty"`
	StackSignature   string    `json:"stackSignature,omitempty"`
	TraceID          string    `json:"traceId,omitempty"`
	SpanKind         uint8     `json:"spanKind,omitempty"`
	PeerIPSource     string    `json:"peerIpSource,omitempty"`
	EdgePeerSource   string    `json:"edgePeerSource,omitempty"`
	EdgeService      string    `json:"edgeService,omitempty"`
	IssueFingerprint string    `json:"issueFingerprint,omitempty"`
	IsRetry          uint8     `json:"isRetry"`
	IsTimeout        uint8     `json:"isTimeout"`
	IsCancellation   uint8     `json:"isCancellation"`
	ExactSignature   string    `json:"exactSignature"`
	PatternSignature string    `json:"patternSignature"`
}

func main() {
	addr := flag.String("addr", envOr("OTEL_CLICKHOUSE_ADDR", "127.0.0.1:9000"), "ClickHouse host:port")
	database := flag.String("database", envOr("OTEL_CLICKHOUSE_DB", "otel"), "ClickHouse database")
	user := flag.String("user", envOr("OTEL_CLICKHOUSE_USER", "default"), "ClickHouse user")
	password := flag.String("password", os.Getenv("OTEL_CLICKHOUSE_PASSWORD"), "ClickHouse password")
	since := flag.String("since", "", "inclusive UTC RFC3339 timestamp")
	until := flag.String("until", "", "exclusive UTC RFC3339 timestamp")
	maxPerPattern := flag.Int("max-per-pattern", 25, "maximum examples per pattern")
	maxEvents := flag.Int("max-events", 10000, "maximum exported events")
	output := flag.String("output", "", "write JSONL to this path instead of stdout")
	flag.Parse()
	if *since == "" || *until == "" || *maxPerPattern < 1 || *maxEvents < 1 {
		log.Fatal("-since, -until, -max-per-pattern, and -max-events are required and must be valid")
	}
	start, err := time.Parse(time.RFC3339Nano, *since)
	if err != nil {
		log.Fatal(err)
	}
	end, err := time.Parse(time.RFC3339Nano, *until)
	if err != nil || !start.Before(end) {
		log.Fatal("invalid training window")
	}

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
	defer cancel()
	conn, err := clickhouse.Open(&clickhouse.Options{Addr: []string{*addr}, Auth: clickhouse.Auth{Database: *database, Username: *user, Password: *password}})
	if err != nil {
		log.Fatal(err)
	}
	defer conn.Close()
	if err := conn.Ping(ctx); err != nil {
		log.Fatal(err)
	}
	writer := os.Stdout
	if *output != "" {
		file, err := os.Create(*output)
		if err != nil {
			log.Fatal(err)
		}
		defer file.Close()
		writer = file
	}
	query := fmt.Sprintf(`SELECT event_id, timestamp, service, release, kind, source, is_error, status_class, status_code, duration_ms,
		route, span_name, metric_name, metric_type, metric_value, metric_count, metric_sum, metric_min, metric_max,
		exception_type, message_normalized, stack_signature, trace_id,
		span_kind, peer_ip_source, edge_peer_source, edge_service, issue_fingerprint, is_retry, is_timeout, is_cancellation, exact_signature, pattern_signature
		FROM %s.normalized_events FINAL
		WHERE timestamp >= toDateTime64('%s', 9, 'UTC') AND timestamp < toDateTime64('%s', 9, 'UTC')
		ORDER BY pattern_signature, timestamp
		LIMIT %d BY pattern_signature
		LIMIT %d`, *database, clickhouseTime(start), clickhouseTime(end), *maxPerPattern, *maxEvents)
	rows, err := conn.Query(ctx, query)
	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()
	encoder := json.NewEncoder(writer)
	count := 0
	for rows.Next() {
		var item sample
		if err := rows.Scan(&item.EventID, &item.Timestamp, &item.Service, &item.Release, &item.Kind, &item.Source, &item.IsError, &item.StatusClass, &item.StatusCode, &item.DurationMS, &item.Route, &item.SpanName, &item.MetricName, &item.MetricType, &item.MetricValue, &item.MetricCount, &item.MetricSum, &item.MetricMin, &item.MetricMax, &item.ExceptionType, &item.Message, &item.StackSignature, &item.TraceID, &item.SpanKind, &item.PeerIPSource, &item.EdgePeerSource, &item.EdgeService, &item.IssueFingerprint, &item.IsRetry, &item.IsTimeout, &item.IsCancellation, &item.ExactSignature, &item.PatternSignature); err != nil {
			log.Fatal(err)
		}
		if err := encoder.Encode(item); err != nil {
			log.Fatal(err)
		}
		count++
	}
	if err := rows.Err(); err != nil {
		log.Fatal(err)
	}
	log.Printf("exported %d training events", count)
}

func clickhouseTime(value time.Time) string {
	return value.UTC().Format("2006-01-02 15:04:05.000000000")
}
func envOr(key, fallback string) string {
	if value := os.Getenv(key); value != "" {
		return value
	}
	return fallback
}