// drive9-billing-worker closes complete UTC hours into the billing usage
// ledger and delivers aggregated usage messages to the billing SQS queue.
// It is a separate binary from drive9-server: it shares the central meta DB
// for the ledger/outbox/snapshot tables and reads per-tenant file ops from
// the billing_ops_delta table (flushed per-pod by drive9-server).
package main

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	awsconfig "github.com/aws/aws-sdk-go-v2/config"

	"github.com/mem9-ai/drive9/pkg/billing"
	"github.com/mem9-ai/drive9/pkg/buildinfo"
	"github.com/mem9-ai/drive9/pkg/leader"
	"github.com/mem9-ai/drive9/pkg/logger"
	"github.com/mem9-ai/drive9/pkg/meta"
	"github.com/mem9-ai/drive9/pkg/metrics"
)

func main() {
	if len(os.Args) > 1 {
		switch os.Args[1] {
		case "version":
			fmt.Println(buildinfo.String("drive9-billing-worker"))
			return
		case "help", "-h", "--help":
			usage(os.Stdout, 0)
		default:
			fmt.Fprintf(os.Stderr, "unknown command %q\n", os.Args[1])
			usage(os.Stderr, 2)
		}
		return
	}
	if err := run(); err != nil {
		fmt.Fprintf(os.Stderr, "drive9-billing-worker: %v\n", err)
		os.Exit(1)
	}
}

func usage(out io.Writer, exitCode int) {
	_, _ = fmt.Fprintf(out, `Usage: drive9-billing-worker [version|help]

Environment:
  DRIVE9_META_DSN                    Central meta DB DSN (required)
  DRIVE9_BILLING_PROVIDER            Deployment provider label (default: aws)
  DRIVE9_BILLING_REGION              Deployment region label (required)
  DRIVE9_BILLING_SQS_QUEUE_URL       Billing report SQS queue URL (required when send enabled)
  DRIVE9_BILLING_SEND_ENABLED        Deliver SQS messages; false = shadow mode, ledger/outbox
                                     only (default: false)
  DRIVE9_BILLING_SNAPSHOT_INTERVAL   Storage snapshot cadence (default: 5m)
  DRIVE9_BILLING_COLLECT_INTERVAL    Hour-collection tick (default: 1m)
  DRIVE9_BILLING_HOUR_CLOSE_DELAY    Wait after an hour ends before closing it (default: 10m)
  DRIVE9_BILLING_SNAPSHOT_STALE_TOLERANCE  Max snapshot age at hour end before the hour is incomplete (default: 2x snapshot interval)
  DRIVE9_BILLING_INCOMPLETE_RETRY_WINDOW   How far back incomplete storage hours are retried (default: 24h)
  DRIVE9_BILLING_SEND_INTERVAL       Outbox send-loop cadence (default: 5s)
  DRIVE9_BILLING_OUTBOX_BATCH_SIZE   Rows claimed per send pass (default: 50)
  DRIVE9_BILLING_OUTBOX_LEASE        Claim lease for in-flight sends (default: 2m)
  DRIVE9_BILLING_MAX_SEND_ATTEMPTS   Attempts before permanent_failed (default: 8)
  DRIVE9_BILLING_LEADER_DISABLED     1/true = always leader (single replica)
  DRIVE9_BILLING_LEADER_LOCK_NAME    GET_LOCK name (default: drive9:billing)
  DRIVE9_BILLING_LISTEN              /metrics + /healthz listen address (default: 127.0.0.1:9091)
`)
	os.Exit(exitCode)
}

func envOr(key, fallback string) string {
	if value := os.Getenv(key); value != "" {
		return value
	}
	return fallback
}

// sqsRegionFromQueueURL extracts the region from a standard SQS queue URL
// (https://sqs.{region}.amazonaws.com/{account}/{queue}); returns "" when the
// host does not match that shape.
func sqsRegionFromQueueURL(queueURL string) string {
	u, err := url.Parse(strings.TrimSpace(queueURL))
	if err != nil {
		return ""
	}
	host := u.Hostname()
	parts := strings.Split(host, ".")
	if len(parts) >= 4 && parts[0] == "sqs" && parts[1] != "" {
		return parts[1]
	}
	return ""
}

func envBool(key string, fallback bool) bool {
	if value := os.Getenv(key); value != "" {
		if parsed, err := strconv.ParseBool(value); err == nil {
			return parsed
		}
		fmt.Fprintf(os.Stderr, "drive9-billing-worker: malformed %s=%q, using default %v\n", key, value, fallback)
	}
	return fallback
}

func envInt(key string, fallback int) int {
	if value := os.Getenv(key); value != "" {
		if parsed, err := strconv.Atoi(value); err == nil {
			return parsed
		}
		fmt.Fprintf(os.Stderr, "drive9-billing-worker: malformed %s=%q, using default %d\n", key, value, fallback)
	}
	return fallback
}

func envDuration(key string, fallback time.Duration) time.Duration {
	if value := os.Getenv(key); value != "" {
		if parsed, err := time.ParseDuration(value); err == nil {
			return parsed
		}
		fmt.Fprintf(os.Stderr, "drive9-billing-worker: malformed %s=%q, using default %v\n", key, value, fallback)
	}
	return fallback
}

func run() error {
	log, err := logger.NewServerLogger()
	if err != nil {
		return err
	}
	defer func() { _ = log.Sync() }()
	ctx := logger.WithContext(context.Background(), log)

	metaDSN := os.Getenv("DRIVE9_META_DSN")
	if metaDSN == "" {
		return fmt.Errorf("DRIVE9_META_DSN is required")
	}
	region := envOr("DRIVE9_BILLING_REGION", "")
	if region == "" {
		return fmt.Errorf("DRIVE9_BILLING_REGION is required")
	}

	cfg := billing.Config{
		Provider:               envOr("DRIVE9_BILLING_PROVIDER", "aws"),
		Region:                 region,
		SQSQueueURL:            envOr("DRIVE9_BILLING_SQS_QUEUE_URL", ""),
		SendEnabled:            envBool("DRIVE9_BILLING_SEND_ENABLED", false),
		SnapshotInterval:       envDuration("DRIVE9_BILLING_SNAPSHOT_INTERVAL", 0),
		CollectInterval:        envDuration("DRIVE9_BILLING_COLLECT_INTERVAL", 0),
		HourCloseDelay:         envDuration("DRIVE9_BILLING_HOUR_CLOSE_DELAY", 0),
		SnapshotStaleTolerance: envDuration("DRIVE9_BILLING_SNAPSHOT_STALE_TOLERANCE", 0),
		IncompleteRetryWindow:  envDuration("DRIVE9_BILLING_INCOMPLETE_RETRY_WINDOW", 0),
		SendInterval:           envDuration("DRIVE9_BILLING_SEND_INTERVAL", 0),
		OutboxBatchSize:        envInt("DRIVE9_BILLING_OUTBOX_BATCH_SIZE", 0),
		OutboxLease:            envDuration("DRIVE9_BILLING_OUTBOX_LEASE", 0),
		MaxSendAttempts:        envInt("DRIVE9_BILLING_MAX_SEND_ATTEMPTS", 0),
	}
	if cfg.SendEnabled && cfg.SQSQueueURL == "" {
		return fmt.Errorf("DRIVE9_BILLING_SQS_QUEUE_URL is required when DRIVE9_BILLING_SEND_ENABLED is true")
	}

	store, err := meta.OpenContext(ctx, metaDSN)
	if err != nil {
		return fmt.Errorf("open meta store: %w", err)
	}
	defer func() { _ = store.Close() }()

	var leaderManager *leader.Manager
	if envBool("DRIVE9_BILLING_LEADER_DISABLED", false) {
		leaderManager = leader.NewManager(store.DB(), leader.WithDisabled())
	} else {
		leaderManager = leader.NewManager(store.DB(),
			leader.WithLockName(envOr("DRIVE9_BILLING_LEADER_LOCK_NAME", "drive9:billing")))
	}
	leaderManager.Start(ctx)
	defer leaderManager.Stop()

	var sqsSender billing.SQSSender
	if cfg.SendEnabled {
		// The queue may live in a different region than the deployment (e.g.
		// billing-report queues are all in us-west-2): derive the SQS client
		// region from the queue URL, falling back to the deployment region.
		sqsRegion := sqsRegionFromQueueURL(cfg.SQSQueueURL)
		if sqsRegion == "" {
			sqsRegion = cfg.Region
		}
		awsCfg, err := awsconfig.LoadDefaultConfig(ctx, awsconfig.WithRegion(sqsRegion))
		if err != nil {
			return fmt.Errorf("load aws config: %w", err)
		}
		sqsSender = billing.NewSQSSender(awsCfg)
	}

	hostname, _ := os.Hostname()
	owner := fmt.Sprintf("%s/%d", hostname, os.Getpid())
	worker := billing.NewWorker(cfg, store, sqsSender, owner, leaderManager)

	metrics.RegisterModule("billing")
	metrics.SetModuleAvailability("billing", true)

	stopCtx, stop := signal.NotifyContext(ctx, os.Interrupt, syscall.SIGTERM)
	defer stop()

	errCh := make(chan error, 2)
	go func() { errCh <- worker.Run(stopCtx) }()
	go func() { errCh <- serveHTTP(stopCtx, envOr("DRIVE9_BILLING_LISTEN", "127.0.0.1:9091")) }()

	err = <-errCh
	if err != nil && stopCtx.Err() == nil {
		return err
	}
	return nil
}

// serveHTTP exposes /metrics (worker + outbox gauges for scraping) and
// /healthz for liveness probes.
func serveHTTP(ctx context.Context, addr string) error {
	mux := http.NewServeMux()
	mux.HandleFunc("/metrics", func(w http.ResponseWriter, r *http.Request) {
		metrics.WritePrometheus(w)
	})
	mux.HandleFunc("/healthz", func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
	})
	srv := &http.Server{Addr: addr, Handler: mux, ReadHeaderTimeout: 10 * time.Second}
	go func() {
		<-ctx.Done()
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		_ = srv.Shutdown(shutdownCtx)
	}()
	if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		return err
	}
	return nil
}