package main

import (
	"bytes"
	"context"
	"database/sql"
	"encoding/csv"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/signal"
	"sort"
	"strings"
	"syscall"
	"time"

	_ "github.com/go-sql-driver/mysql"
)

func main() {
	config, err := mainConfig()
	if err != nil {
		log.Fatal(err)
	}

	db, err := sql.Open("mysql", config.DBDSN)
	if err != nil {
		log.Fatal("open database: ", err)
	}
	defer db.Close()

	httpClient := &http.Client{Timeout: config.RequestTimeout}
	store := Store{db: db}
	notifier := Notifier{webhook: config.FeishuWebhook, httpClient: httpClient}
	collector := Collector{
		client:   O11YClient{httpClient: httpClient, config: config},
		store:    store,
		notifier: notifier,
		timeout:  config.RequestTimeout,
	}

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	log.Printf("collector started; poll interval=%s", config.PollInterval)
	if err := collector.runAndNotify(ctx, time.Now()); err != nil {
		log.Printf("initial collection failed: %v", err)
	}

	nextScheduledAt := nextScheduledRun(time.Now())
	for {
		waitDuration := time.Until(nextScheduledAt)
		if waitDuration < 0 {
			waitDuration = 0
		}
		timer := time.NewTimer(waitDuration)
		select {
		case <-ctx.Done():
			timer.Stop()
			log.Println("collector stopped")
			return
		case <-timer.C:
			log.Printf("scheduled collection triggered at %s", nextScheduledAt.UTC().Format(time.RFC3339))
			if err := collector.runAndNotify(ctx, nextScheduledAt); err != nil {
				log.Printf("collection failed: %v", err)
			}
			nextScheduledAt = nextScheduledRun(nextScheduledAt)
		}
	}
}

type Collector struct {
	client   O11YClient
	store    Store
	notifier Notifier
	timeout  time.Duration
	now      func() time.Time
}

func (c Collector) runAndNotify(parent context.Context, scheduledAt time.Time) error {
	err := c.run(parent, scheduledAt)
	if err == nil {
		return nil
	}

	notificationCtx, cancel := context.WithTimeout(context.Background(), c.timeout)
	defer cancel()
	if notifyErr := c.notifier.notifyFailure(notificationCtx, err); notifyErr != nil {
		log.Printf("send failure notification: %v", notifyErr)
	}
	return err
}

func (c Collector) run(parent context.Context, scheduledAt time.Time) (runErr error) {
	ctx, cancel := context.WithTimeout(parent, c.timeout)
	defer cancel()

	log.Printf("collection started; scheduled_at=%s", scheduledAt.UTC().Format(time.RFC3339))
	if err := c.store.ensureSchema(ctx); err != nil {
		return err
	}
	if err := c.notifyMissedRun(ctx, scheduledAt); err != nil {
		log.Printf("detect missed collection: %v", err)
	}
	runID, err := c.store.startTaskRun(ctx, scheduledAt)
	if err != nil {
		return fmt.Errorf("record task start: %w", err)
	}
	log.Printf("task run recorded; id=%d", runID)
	clusterCount := 0
	defer func() {
		finishCtx, finishCancel := context.WithTimeout(context.Background(), c.timeout)
		defer finishCancel()
		status, message := "succeeded", ""
		if runErr != nil {
			status, message = "failed", runErr.Error()
		}
		if err := c.store.finishTaskRun(finishCtx, runID, status, clusterCount, message); err != nil {
			log.Printf("record task completion; id=%d error=%v", runID, err)
			if runErr == nil {
				runErr = fmt.Errorf("record task completion: %w", err)
			}
			return
		}
		log.Printf("task run finished; id=%d status=%s cluster_count=%d", runID, status, clusterCount)
	}()

	log.Printf("authenticating with O11Y")
	token, err := c.client.authenticate(ctx)
	if err != nil {
		return fmt.Errorf("authenticate: %w", err)
	}
	log.Printf("O11Y authentication succeeded")
	now := time.Now()
	if c.now != nil {
		now = c.now()
	}
	urls, err := c.client.queryExport(ctx, token, now)
	if err != nil {
		return fmt.Errorf("query OLAP: %w", err)
	}
	log.Printf("OLAP export succeeded; file_count=%d", len(urls))

	uniqueClusters := make(map[string]struct{})
	for index, url := range urls {
		log.Printf("downloading OLAP export; file=%d/%d", index+1, len(urls))
		clusterIDs, err := c.client.downloadClusters(ctx, url)
		if err != nil {
			return fmt.Errorf("download export: %w", err)
		}
		log.Printf("parsed OLAP export; file=%d cluster_count=%d", index+1, len(clusterIDs))
		for _, clusterID := range clusterIDs {
			uniqueClusters[clusterID] = struct{}{}
		}
	}
	clusterIDs := make([]string, 0, len(uniqueClusters))
	for clusterID := range uniqueClusters {
		clusterIDs = append(clusterIDs, clusterID)
	}
	log.Printf("upserting active cluster IDs; count=%d", len(clusterIDs))
	if err := c.store.upsertClusters(ctx, clusterIDs); err != nil {
		return fmt.Errorf("store cluster IDs: %w", err)
	}
	clusterCount = len(clusterIDs)
	log.Printf("collection completed: %d cluster IDs", len(clusterIDs))
	return nil
}

func (c Collector) notifyMissedRun(ctx context.Context, scheduledAt time.Time) error {
	lastStartedAt, found, err := c.store.lastTaskRunStartedAt(ctx)
	if err != nil || !found || !isMissedRun(lastStartedAt, scheduledAt) {
		return err
	}
	missedDuration := scheduledAt.Sub(lastStartedAt).Round(time.Second)
	err = fmt.Errorf("collection gap detected: last started at %s, gap=%s",
		lastStartedAt.UTC().Format(time.RFC3339), missedDuration)
	log.Printf("%v", err)
	notificationCtx, cancel := context.WithTimeout(context.Background(), c.timeout)
	defer cancel()
	if notifyErr := c.notifier.notifyFailure(notificationCtx, err); notifyErr != nil {
		return fmt.Errorf("send missed-run notification: %w", notifyErr)
	}
	return nil
}

func isMissedRun(lastStartedAt, scheduledAt time.Time) bool {
	return scheduledAt.Sub(lastStartedAt) >= missedRunThreshold
}

func nextScheduledRun(after time.Time) time.Time {
	// A completed task is passed back as an exact schedule boundary. Start
	// from the following minute so that boundary cannot be selected again.
	next := after.Truncate(time.Minute).Add(time.Minute)

	for {
		minute := next.Minute()
		if minute%10 == 5 {
			return next
		}
		next = next.Add(time.Minute)
	}
}

const (
	authURL            = "https://www.ds.us-west-2.aws.observability.tidbcloud.com/client/auth"
	olapURL            = "https://www.ds.us-east-1-serverless.aws.observability.tidbcloud.com/internal/o11y/diagnosis/api/v1/olap"
	o11yProvider       = "aws"
	o11yRegion         = "us-east-1"
	olapOrgID          = "1372813089209260919"
	clusterTable       = "manus.active_cluster_ids"
	taskRunTable       = "manus.active_cluster_task_runs"
	pollInterval       = 10 * time.Minute
	queryWindow        = 10 * time.Minute
	requestTimeout     = time.Minute
	missedRunThreshold = 2 * pollInterval
)

type Config struct {
	GitHubToken    string
	DBDSN          string
	FeishuWebhook  string
	AuthURL        string
	OLAPURL        string
	Provider       string
	Region         string
	PollInterval   time.Duration
	QueryWindow    time.Duration
	RequestTimeout time.Duration
}

func loadConfig(getenv func(string) string) (Config, error) {
	config := Config{
		GitHubToken:    strings.TrimSpace(getenv("GITHUB_TOKEN")),
		DBDSN:          strings.TrimSpace(getenv("DB_DSN")),
		FeishuWebhook:  strings.TrimSpace(getenv("FEISHU_WEBHOOK")),
		AuthURL:        authURL,
		OLAPURL:        olapURL,
		Provider:       o11yProvider,
		Region:         o11yRegion,
		PollInterval:   pollInterval,
		QueryWindow:    queryWindow,
		RequestTimeout: requestTimeout,
	}
	for name, value := range map[string]string{
		"GITHUB_TOKEN": config.GitHubToken,
		"DB_DSN":       config.DBDSN,
	} {
		if value == "" {
			return Config{}, fmt.Errorf("%s is required", name)
		}
	}
	return config, nil
}

func mainConfig() (Config, error) {
	return loadConfig(os.Getenv)
}

type O11YClient struct {
	httpClient *http.Client
	config     Config
}

type authResponse struct {
	Token string `json:"token"`
}

type olapResponse struct {
	Status string `json:"status"`
	Data   struct {
		Status string `json:"status"`
		Result struct {
			Status       string `json:"status"`
			DownloadURLs []struct {
				URL string `json:"url"`
			} `json:"download_urls"`
		} `json:"result"`
	} `json:"data"`
}

func (c O11YClient) authenticate(ctx context.Context) (string, error) {
	body, err := json.Marshal(map[string]string{
		"github_token": c.config.GitHubToken,
		"provider":     c.config.Provider,
	})
	if err != nil {
		return "", fmt.Errorf("encode auth request: %w", err)
	}
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, c.config.AuthURL, bytes.NewReader(body))
	if err != nil {
		return "", fmt.Errorf("create auth request: %w", err)
	}
	request.Header.Set("Content-Type", "application/json")
	response, err := c.httpClient.Do(request)
	if err != nil {
		return "", fmt.Errorf("send auth request: %w", err)
	}
	defer response.Body.Close()
	if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
		return "", responseError("auth", response)
	}
	var payload authResponse
	if err := json.NewDecoder(response.Body).Decode(&payload); err != nil {
		return "", fmt.Errorf("decode auth response: %w", err)
	}
	if payload.Token == "" {
		return "", fmt.Errorf("auth response has no token")
	}
	return payload.Token, nil
}

func (c O11YClient) queryExport(ctx context.Context, token string, now time.Time) ([]string, error) {
	date := now.UTC().Format("2006-01-02")
	table := "sqlstatement_org_" + olapOrgID + "_all"
	sql := fmt.Sprintf(
		`SELECT DISTINCT serverless_cluster_id FROM %s WHERE "date" >= '%s' AND "end" >= to_unixtime(current_timestamp) - %d`,
		table, date, int(c.config.QueryWindow.Seconds()),
	)
	body, err := json.Marshal(map[string]any{
		"sql":            sql,
		"output_mode":    "export",
		"inject_cluster": false,
		"inject_tenant":  false,
	})
	if err != nil {
		return nil, fmt.Errorf("encode OLAP query: %w", err)
	}
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, c.config.OLAPURL, bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("create OLAP request: %w", err)
	}
	request.Header.Set("Accept", "application/json, text/plain, */*")
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("x-cluster-id", "all")
	request.Header.Set("x-provider", c.config.Provider)
	request.Header.Set("x-region", c.config.Region)
	request.Header.Set("o11ytoken", token)
	response, err := c.httpClient.Do(request)
	if err != nil {
		return nil, fmt.Errorf("send OLAP request: %w", err)
	}
	defer response.Body.Close()
	if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
		return nil, responseError("OLAP", response)
	}
	var payload olapResponse
	if err := json.NewDecoder(response.Body).Decode(&payload); err != nil {
		return nil, fmt.Errorf("decode OLAP response: %w", err)
	}
	if payload.Status != "ok" || payload.Data.Status != "SUCCEEDED" || payload.Data.Result.Status != "SUCCEEDED" {
		return nil, fmt.Errorf("OLAP export did not succeed: status=%q data_status=%q result_status=%q",
			payload.Status, payload.Data.Status, payload.Data.Result.Status)
	}
	urls := make([]string, 0, len(payload.Data.Result.DownloadURLs))
	for _, download := range payload.Data.Result.DownloadURLs {
		if download.URL != "" {
			urls = append(urls, download.URL)
		}
	}
	if len(urls) == 0 {
		return nil, fmt.Errorf("OLAP export returned no download URLs")
	}
	return urls, nil
}

func (c O11YClient) downloadClusters(ctx context.Context, url string) ([]string, error) {
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, fmt.Errorf("create download request: %w", err)
	}
	response, err := c.httpClient.Do(request)
	if err != nil {
		return nil, fmt.Errorf("download CSV: %w", err)
	}
	defer response.Body.Close()
	if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
		return nil, responseError("download", response)
	}
	return parseClusterCSV(response.Body)
}

func responseError(stage string, response *http.Response) error {
	limitedBody, _ := io.ReadAll(io.LimitReader(response.Body, 4<<10))
	return fmt.Errorf("%s response status %s: %q", stage, response.Status, string(limitedBody))
}

func parseClusterCSV(input io.Reader) ([]string, error) {
	reader := csv.NewReader(input)
	headers, err := reader.Read()
	if err != nil {
		return nil, fmt.Errorf("read CSV header: %w", err)
	}
	clusterColumn := -1
	for index, header := range headers {
		switch strings.TrimSpace(strings.TrimPrefix(header, "\ufeff")) {
		case "serverless_cluster_id", "cluster_id":
			clusterColumn = index
		}
	}
	if clusterColumn == -1 {
		return nil, fmt.Errorf("CSV does not contain serverless_cluster_id or cluster_id column")
	}
	unique := make(map[string]struct{})
	for {
		record, err := reader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, fmt.Errorf("read CSV row: %w", err)
		}
		if len(record) <= clusterColumn {
			return nil, fmt.Errorf("CSV row has %d columns; expected cluster ID column %d", len(record), clusterColumn)
		}
		if clusterID := strings.TrimSpace(record[clusterColumn]); clusterID != "" {
			unique[clusterID] = struct{}{}
		}
	}
	clusters := make([]string, 0, len(unique))
	for clusterID := range unique {
		clusters = append(clusters, clusterID)
	}
	sort.Strings(clusters)
	return clusters, nil
}

type Store struct {
	db *sql.DB
}

func (s Store) ensureSchema(ctx context.Context) error {
	if _, err := s.db.ExecContext(ctx, "CREATE DATABASE IF NOT EXISTS manus"); err != nil {
		return fmt.Errorf("create manus database: %w", err)
	}
	log.Printf("database manus is ready")
	if _, err := s.db.ExecContext(ctx, fmt.Sprintf(`
CREATE TABLE IF NOT EXISTS %s (
  cluster_id VARCHAR(255) NOT NULL,
  created_at DATETIME(6) NOT NULL,
  updated_at DATETIME(6) NOT NULL,
  update_count BIGINT UNSIGNED NOT NULL DEFAULT 1,
  PRIMARY KEY (cluster_id)
)`, clusterTable)); err != nil {
		return fmt.Errorf("create active_cluster_ids table: %w", err)
	}
	if _, err := s.db.ExecContext(ctx, fmt.Sprintf(`
ALTER TABLE %s
ADD COLUMN IF NOT EXISTS update_count BIGINT UNSIGNED NOT NULL DEFAULT 1`, clusterTable)); err != nil {
		return fmt.Errorf("add update_count column: %w", err)
	}
	if _, err := s.db.ExecContext(ctx, fmt.Sprintf(`
CREATE TABLE IF NOT EXISTS %s (
  id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,
  scheduled_at DATETIME(6) NOT NULL,
  started_at DATETIME(6) NOT NULL,
  finished_at DATETIME(6) NULL,
  status VARCHAR(16) NOT NULL,
  cluster_count INT UNSIGNED NOT NULL DEFAULT 0,
  error_message TEXT NULL,
  PRIMARY KEY (id),
  KEY idx_active_cluster_task_runs_started_at (started_at)
)`, taskRunTable)); err != nil {
		return fmt.Errorf("create active_cluster_task_runs table: %w", err)
	}
	log.Printf("database schema is ready")
	return nil
}

func (s Store) lastTaskRunStartedAt(ctx context.Context) (time.Time, bool, error) {
	var startedAt time.Time
	err := s.db.QueryRowContext(ctx, fmt.Sprintf(`
SELECT started_at
FROM %s
ORDER BY started_at DESC
LIMIT 1`, taskRunTable)).Scan(&startedAt)
	if err == sql.ErrNoRows {
		return time.Time{}, false, nil
	}
	if err != nil {
		return time.Time{}, false, fmt.Errorf("query latest task run: %w", err)
	}
	return startedAt, true, nil
}

func (s Store) startTaskRun(ctx context.Context, scheduledAt time.Time) (int64, error) {
	result, err := s.db.ExecContext(ctx, fmt.Sprintf(`
INSERT INTO %s (scheduled_at, started_at, status)
VALUES (?, CURRENT_TIMESTAMP(6), 'running')`, taskRunTable), scheduledAt.UTC())
	if err != nil {
		return 0, fmt.Errorf("insert task run: %w", err)
	}
	runID, err := result.LastInsertId()
	if err != nil {
		return 0, fmt.Errorf("get task run ID: %w", err)
	}
	return runID, nil
}

func (s Store) finishTaskRun(ctx context.Context, runID int64, status string, clusterCount int, errorMessage string) error {
	_, err := s.db.ExecContext(ctx, fmt.Sprintf(`
UPDATE %s
SET finished_at = CURRENT_TIMESTAMP(6), status = ?, cluster_count = ?, error_message = ?
WHERE id = ?`, taskRunTable), status, clusterCount, nullableString(errorMessage), runID)
	if err != nil {
		return fmt.Errorf("update task run: %w", err)
	}
	return nil
}

func nullableString(value string) any {
	if value == "" {
		return nil
	}
	return value
}

func (s Store) upsertClusters(ctx context.Context, clusterIDs []string) error {
	if len(clusterIDs) == 0 {
		return nil
	}
	const batchSize = 500
	for start := 0; start < len(clusterIDs); start += batchSize {
		end := start + batchSize
		if end > len(clusterIDs) {
			end = len(clusterIDs)
		}
		if err := s.upsertBatch(ctx, clusterIDs[start:end]); err != nil {
			return err
		}
	}
	return nil
}

func (s Store) upsertBatch(ctx context.Context, clusterIDs []string) error {
	query, args := buildUpsertBatch(clusterIDs)
	if _, err := s.db.ExecContext(ctx, query, args...); err != nil {
		return fmt.Errorf("upsert cluster IDs: %w", err)
	}
	return nil
}

func buildUpsertBatch(clusterIDs []string) (string, []any) {
	values := make([]string, 0, len(clusterIDs))
	args := make([]any, 0, len(clusterIDs))
	for _, clusterID := range clusterIDs {
		values = append(values, "(?, CURRENT_TIMESTAMP(6), CURRENT_TIMESTAMP(6), 1)")
		args = append(args, clusterID)
	}
	query := fmt.Sprintf(`
INSERT INTO %s (cluster_id, created_at, updated_at, update_count)
VALUES %s
ON DUPLICATE KEY UPDATE updated_at = CURRENT_TIMESTAMP(6), update_count = update_count + 1`, clusterTable, strings.Join(values, ","))
	return query, args
}

type Notifier struct {
	webhook    string
	httpClient *http.Client
}

func (n Notifier) notifyFailure(ctx context.Context, err error) error {
	if n.webhook == "" {
		return nil
	}
	body, marshalErr := json.Marshal(map[string]any{
		"msg_type": "text",
		"content": map[string]string{
			"text": "o11y-olap collector failed: " + err.Error(),
		},
	})
	if marshalErr != nil {
		return fmt.Errorf("encode Feishu notification: %w", marshalErr)
	}
	request, requestErr := http.NewRequestWithContext(ctx, http.MethodPost, n.webhook, bytes.NewReader(body))
	if requestErr != nil {
		return fmt.Errorf("create Feishu request: %w", requestErr)
	}
	request.Header.Set("Content-Type", "application/json")
	response, doErr := n.httpClient.Do(request)
	if doErr != nil {
		return fmt.Errorf("send Feishu notification: %w", doErr)
	}
	defer response.Body.Close()
	if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
		return responseError("Feishu", response)
	}
	return nil
}