package main

import (
	"context"
	"errors"
	"log/slog"
	"net"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"git.yuntuo.hk/k2safe/orbit/internal/config"
	"git.yuntuo.hk/k2safe/orbit/internal/credentialvault"
	"git.yuntuo.hk/k2safe/orbit/internal/observability"
	"git.yuntuo.hk/k2safe/orbit/internal/retention"
	"git.yuntuo.hk/k2safe/orbit/internal/storage/mysql"
	"git.yuntuo.hk/k2safe/orbit/internal/transport/httpserver"
	core "git.yuntuo.hk/k2safe/orbit/model"
	"git.yuntuo.hk/k2safe/orbit/platform/events"
	"git.yuntuo.hk/k2safe/orbit/platform/resources"
	"git.yuntuo.hk/k2safe/orbit/pulse/schedule"
	"git.yuntuo.hk/k2safe/orbit/pulse/service"
	"git.yuntuo.hk/k2safe/orbit/pulse/worker"
	"git.yuntuo.hk/k2safe/orbit/weave/control"
	"git.yuntuo.hk/k2safe/orbit/weave/definition"
	"git.yuntuo.hk/k2safe/orbit/weave/instance"
	"git.yuntuo.hk/k2safe/orbit/weave/recovery"
	weavesignal "git.yuntuo.hk/k2safe/orbit/weave/signal"
)

func main() {
	if err := run(); err != nil {
		slog.Error("orbit stopped", "error", err)
		os.Exit(1)
	}
}
func run() error {
	slog.SetDefault(slog.New(slog.NewJSONHandler(os.Stdout, nil)))
	cfg, err := config.Load()
	if err != nil {
		return err
	}
	vault, err := credentialvault.LoadFile(cfg.CredentialKeyFile)
	if err != nil {
		return err
	}
	if cfg.ConsoleApplicationAdmin != "" && vault == nil {
		return errors.New("ORBIT_CREDENTIAL_KEY_FILE is required for console application administration")
	}
	root, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	startup, cancel := context.WithTimeout(root, 10*time.Second)
	db, err := mysql.Open(startup, cfg.MySQLDSN, cfg.MaxOpen, cfg.MaxIdle)
	if err != nil {
		cancel()
		return err
	}
	defer db.Close()
	if err := mysql.CheckSchema(startup, db); err != nil {
		cancel()
		return err
	}
	cancel()
	readiness := httpserver.NewReadiness(func(ctx context.Context) error {
		if err := db.PingContext(ctx); err != nil {
			return err
		}
		return mysql.CheckSchema(ctx, db)
	})
	var metrics *observability.Registry
	var metricsListener net.Listener
	if cfg.MetricsAddr != "" {
		metricsListener, err = net.Listen("tcp", cfg.MetricsAddr)
		if err != nil {
			return errors.New("metrics listener could not bind")
		}
		defer metricsListener.Close()
		metrics = observability.NewRegistry()
		metrics.RecordRuntimeConfig(observability.RuntimeConfig{
			WorkflowTickInterval:  cfg.WorkflowTickInterval,
			ScheduleTickInterval:  cfg.ScheduleTickInterval,
			RecoveryTickInterval:  cfg.RecoveryTickInterval,
			BatchLimit:            cfg.RuntimeBatchLimit,
			ShutdownTimeout:       cfg.ShutdownTimeout,
			WorkerConnectionLimit: cfg.WorkerConnectionLimit,
			EventConnectionLimit:  cfg.EventConnectionLimit,
			DispatchWorkers:       cfg.DispatchWorkers,
			DispatchCapacity:      cfg.DispatchCapacity,
		})
		metrics.RecordDatabase(db.Stats())
		txStats := mysql.ReadTransactionStats()
		metrics.RecordTransactionStats(observability.TransactionStats{Retries: txStats.Retries, Exhausted: txStats.Exhausted})
	}
	consoleAuth := mysql.NewConsoleAuthStore(db)
	retentionStore := mysql.NewRetentionPolicyStore(db, cfg.RetentionDefaults)
	server := httpserver.NewWithConsoleAuth(cfg.HTTPAddr, readiness.Check, consoleAuth, service.New(mysql.NewTaskStore(db)))
	httpserver.RegisterApplicationResources(server, resources.New(mysql.NewTaskStore(db)))
	node, err := core.NewID()
	if err != nil {
		return err
	}
	executions, err := mysql.NewExecutionStore(db, node, mysql.ExecutionOptions{})
	if err != nil {
		return err
	}
	hub := httpserver.NewWorkerHub(root, worker.New(executions), httpserver.HubOptions{
		ConnectionLimit:  cfg.WorkerConnectionLimit,
		DispatchWorkers:  cfg.DispatchWorkers,
		DispatchCapacity: cfg.DispatchCapacity,
	})
	httpserver.RegisterWorkers(server, hub)
	eventService := events.New(mysql.NewEventStore(db))
	eventHub := httpserver.NewEventHub(root, eventService, httpserver.HubOptions{
		ConnectionLimit:  cfg.EventConnectionLimit,
		DispatchWorkers:  cfg.DispatchWorkers,
		DispatchCapacity: cfg.DispatchCapacity,
	})
	httpserver.RegisterEventSubscriptions(server, eventService)
	httpserver.RegisterEventConsumers(server, eventHub)
	defer func() {
		ctx, cancel := context.WithTimeout(context.Background(), cfg.ShutdownTimeout)
		defer cancel()
		_ = eventHub.Shutdown(ctx)
	}()
	schedules, err := mysql.NewScheduleStore(db, node)
	if err != nil {
		return err
	}
	httpserver.RegisterSchedules(server, schedule.New(schedules))
	httpserver.RegisterWorkflowDefinitions(server, definition.New(mysql.NewWorkflowDefinitionStore(db)))
	httpserver.RegisterWorkflows(server, instance.New(mysql.NewWorkflowInstanceStore(db)), control.New(mysql.NewWorkflowControlStore(db)), weavesignal.New(mysql.NewWorkflowSignalStore(db)), recovery.New(mysql.NewWorkflowRecoveryStore(db)))
	httpserver.RegisterConsoleRetention(server, consoleAuth, retentionStore, retentionStore)
	httpserver.RegisterConsoleArchiveJobs(server, consoleAuth, retentionStore, cfg.ArchiveStorageID)
	consoleApplications := mysql.NewConsoleApplicationStore(db, cfg.ConsoleApplicationAdmin, vault)
	httpserver.RegisterConsoleApplications(server, consoleAuth, consoleApplications)
	httpserver.RegisterConsoleScopes(server, consoleApplications)
	httpserver.WrapConsoleAuth(server, consoleAuth, consoleApplications)
	httpserver.WrapMetrics(server, metrics)

	if cfg.ArchiveRoot != "" {
		archives := &retention.ArchiveRuntime{Directory: cfg.ArchiveRoot, StorageID: cfg.ArchiveStorageID, MaxBytes: cfg.ArchiveMaxBytes, Store: retentionStore}
		if err := archives.Init(); err != nil {
			return errors.New("archive volume initialization failed; check path, ownership and storage identity")
		}
		archiveCtx, stopArchives := context.WithCancel(root)
		archiveDone := make(chan struct{})
		go func() {
			defer close(archiveDone)
			ticker := time.NewTicker(cfg.ArchiveTickInterval)
			defer ticker.Stop()
			for {
				select {
				case <-archiveCtx.Done():
					return
				case <-ticker.C:
					call, cancel := context.WithTimeout(archiveCtx, 45*time.Second)
					err := archives.Tick(call)
					cancel()
					if err != nil && archiveCtx.Err() == nil {
						slog.Error("archive batch could not be confirmed", "node_id", node)
					}
				}
			}
		}()
		defer func() { stopArchives(); <-archiveDone }()
	}
	workflows, err := mysql.NewWorkflowRuntime(db, node)
	if err != nil {
		return err
	}
	workflowCtx, stopWorkflows := context.WithCancel(root)
	workflowDone := make(chan struct{})
	go func() {
		defer close(workflowDone)
		ticker := time.NewTicker(cfg.WorkflowTickInterval)
		defer ticker.Stop()
		for {
			select {
			case <-workflowCtx.Done():
				return
			case <-ticker.C:
				started := time.Now()
				call, cancel := context.WithTimeout(workflowCtx, 5*time.Second)
				_, err := workflows.Tick(call, cfg.RuntimeBatchLimit)
				cancel()
				if metrics != nil {
					metrics.RecordBatch("workflow", started, err)
				}
				if err != nil && workflowCtx.Err() == nil {
					slog.Error("workflow batch failed", "node_id", node)
				}
			}
		}
	}()
	defer func() { stopWorkflows(); <-workflowDone }()
	scheduleCtx, stopSchedules := context.WithCancel(root)
	scheduleDone := make(chan struct{})
	go func() {
		defer close(scheduleDone)
		ticker := time.NewTicker(cfg.ScheduleTickInterval)
		defer ticker.Stop()
		for {
			select {
			case <-scheduleCtx.Done():
				return
			case <-ticker.C:
				started := time.Now()
				call, cancel := context.WithTimeout(scheduleCtx, 5*time.Second)
				_, err := schedules.Tick(call, cfg.RuntimeBatchLimit)
				cancel()
				if metrics != nil {
					metrics.RecordBatch("schedule", started, err)
				}
				if err != nil && scheduleCtx.Err() == nil {
					slog.Error("schedule batch failed", "node_id", node)
				}
			}
		}
	}()
	defer func() { stopSchedules(); <-scheduleDone }()
	defer func() {
		ctx, cancel := context.WithTimeout(context.Background(), cfg.ShutdownTimeout)
		defer cancel()
		_ = hub.Shutdown(ctx)
	}()
	recoveryDone := make(chan struct{})
	recoveryCtx, stopRecovery := context.WithCancel(root)
	go func() {
		defer close(recoveryDone)
		ticker := time.NewTicker(cfg.RecoveryTickInterval)
		defer ticker.Stop()
		for {
			select {
			case <-recoveryCtx.Done():
				return
			case <-ticker.C:
				started := time.Now()
				call, cancel := context.WithTimeout(recoveryCtx, 5*time.Second)
				_, err := executions.Recover(call, cfg.RuntimeBatchLimit)
				cancel()
				if metrics != nil {
					metrics.RecordBatch("execution_recovery", started, err)
				}
				if err != nil && recoveryCtx.Err() == nil {
					slog.Error("execution recovery batch failed", "node_id", node)
				}
			}
		}
	}()
	defer func() { stopRecovery(); <-recoveryDone }()
	if metrics != nil {
		metricsCtx, stopMetrics := context.WithCancel(root)
		metricsDone := make(chan struct{})
		go func() {
			defer close(metricsDone)
			ticker := time.NewTicker(5 * time.Second)
			defer ticker.Stop()
			for {
				select {
				case <-metricsCtx.Done():
					return
				case <-ticker.C:
					metrics.RecordDatabase(db.Stats())
					txStats := mysql.ReadTransactionStats()
					metrics.RecordTransactionStats(observability.TransactionStats{Retries: txStats.Retries, Exhausted: txStats.Exhausted})
					metrics.RecordProcess()
					started := time.Now()
					call, cancel := context.WithTimeout(metricsCtx, 900*time.Millisecond)
					snapshot, err := mysql.ReadOperationalSnapshot(call, db)
					cancel()
					metrics.RecordBusinessSnapshot(observability.BusinessSnapshot{
						TaskStates: snapshot.TaskStates, TaskOldestAgeSeconds: snapshot.TaskOldestAgeSeconds,
						DeliveryStates: snapshot.DeliveryStates, DeliveryOldestAgeSeconds: snapshot.DeliveryOldestAgeSeconds,
						WorkflowStates: snapshot.WorkflowStates, WorkflowOldestAgeSeconds: snapshot.WorkflowOldestAgeSeconds,
						InboxStates: snapshot.InboxStates, InboxOldestAgeSeconds: snapshot.InboxOldestAgeSeconds,
						RecentEvents:               snapshot.RecentEvents,
						TaskTerminalLatencySeconds: snapshot.TaskTerminalLatencySeconds, TaskTerminalLatencyMaxSeconds: snapshot.TaskTerminalLatencyMaxSeconds,
						WorkflowTerminalLatencySeconds: snapshot.WorkflowTerminalLatencySeconds, WorkflowTerminalLatencyMaxSeconds: snapshot.WorkflowTerminalLatencyMaxSeconds,
						TaskAcceptToStartLatencySeconds: snapshot.TaskAcceptToStartLatencySeconds, TaskAcceptToStartLatencyMaxSeconds: snapshot.TaskAcceptToStartLatencyMaxSeconds,
						TaskAcceptToTerminalLatencySeconds: snapshot.TaskAcceptToTerminalLatencySeconds, TaskAcceptToTerminalLatencyMaxSeconds: snapshot.TaskAcceptToTerminalLatencyMaxSeconds,
						WorkflowStartToTerminalLatencySeconds: snapshot.WorkflowStartToTerminalLatencySeconds, WorkflowStartToTerminalLatencyMaxSeconds: snapshot.WorkflowStartToTerminalLatencyMaxSeconds,
						Truncated: snapshot.Truncated,
					}, started, err)
					if err != nil && metricsCtx.Err() == nil {
						reason := "query_error"
						if errors.Is(err, context.DeadlineExceeded) {
							reason = "timeout"
						} else if errors.Is(err, context.Canceled) {
							reason = "cancelled"
						}
						slog.Error("business metrics snapshot failed", "node_id", node, "reason", reason)
					}
				}
			}
		}()
		defer func() { stopMetrics(); <-metricsDone }()
	}
	listener, err := net.Listen("tcp", cfg.HTTPAddr)
	if err != nil {
		return err
	}
	defer server.Close()
	var metricsServer *http.Server
	done := make(chan error, 2)
	if metrics != nil {
		metricsServer = httpserver.NewMetricsServer(cfg.MetricsAddr, metrics)
		defer metricsServer.Close()
		go func() { done <- metricsServer.Serve(metricsListener) }()
		slog.Info("Orbit metrics server starting", "metrics_address", metricsListener.Addr().String())
	}
	go func() { done <- server.Serve(listener) }()
	slog.Info("Orbit HTTP server starting", "address", listener.Addr().String(), "scope", "health, task, queue, schedule, subscription, workflow and Worker APIs", "node_id", node, "workflow_tick_interval", cfg.WorkflowTickInterval.String(), "schedule_tick_interval", cfg.ScheduleTickInterval.String(), "recovery_tick_interval", cfg.RecoveryTickInterval.String(), "runtime_batch_limit", cfg.RuntimeBatchLimit, "shutdown_timeout", cfg.ShutdownTimeout.String(), "worker_connection_limit", cfg.WorkerConnectionLimit, "event_connection_limit", cfg.EventConnectionLimit, "dispatch_workers", cfg.DispatchWorkers, "dispatch_capacity", cfg.DispatchCapacity)
	select {
	case err := <-done:
		if errors.Is(err, http.ErrServerClosed) {
			return nil
		}
		return err
	case <-root.Done():
		readiness.BeginDrain()
		if metrics != nil {
			metrics.RecordDraining(true)
		}
		shutdown, cancel := context.WithTimeout(context.Background(), cfg.ShutdownTimeout)
		defer cancel()
		if err := server.Shutdown(shutdown); err != nil {
			_ = server.Close()
			return err
		}
		if metricsServer != nil {
			if err := metricsServer.Shutdown(shutdown); err != nil {
				return err
			}
		}
		count := 1
		if metricsServer != nil {
			count++
		}
		for i := 0; i < count; i++ {
			if err := <-done; !errors.Is(err, http.ErrServerClosed) {
				return err
			}
		}
		return nil
	}
}