package main

import (
	"compress/gzip"
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/robfig/cron/v3"
)

// Config holds environmental configurations
type Config struct {
	Host         string
	Port         string
	User         string
	Password     string
	DB           string
	BackupDir    string
	KeepCount    int
	Schedule     string
	AWSAccessKey string
	AWSSecretKey string
	S3Bucket     string
	S3Endpoint   string
	AWSRegion    string
}

func loadConfig() Config {
	keepCount, _ := strconv.Atoi(getEnv("BACKUP_KEEP_COUNT", "7"))
	return Config{
		Host:         getEnv("POSTGRES_HOST", "db"),
		Port:         getEnv("POSTGRES_PORT", "5432"),
		User:         getEnv("POSTGRES_USER", "postgres"),
		Password:     getEnv("POSTGRES_PASSWORD", ""),
		DB:           getEnv("POSTGRES_DB", ""),
		BackupDir:    getEnv("BACKUP_DIR", "/backups"),
		KeepCount:    keepCount,
		Schedule:     os.Getenv("BACKUP_SCHEDULE"),
		AWSAccessKey: os.Getenv("AWS_ACCESS_KEY_ID"),
		AWSSecretKey: os.Getenv("AWS_SECRET_ACCESS_KEY"),
		S3Bucket:     os.Getenv("S3_BUCKET"),
		S3Endpoint:   os.Getenv("S3_ENDPOINT"),
		AWSRegion:    getEnv("AWS_DEFAULT_REGION", "us-east-1"),
	}
}

func main() {
	if len(os.Args) < 2 || os.Args[1] == "help" || os.Args[1] == "-h" || os.Args[1] == "--help" {
		printHelp()
		os.Exit(0)
	}

	cfg := loadConfig()
	mode := os.Args[1]

	switch mode {
	case "cron-init":
		initScheduler(cfg)
	case "run":
		if err := runBackupPipeline(cfg); err != nil {
			logError("Pipeline failed: %v", err)
			os.Exit(1)
		}
	case "inspect":
		if len(os.Args) < 3 {
			logError("ERROR: No backup filename/S3 URL provided!")
			fmt.Println("Usage: backup-app inspect <filename.sql.gz | s3://bucket/key>")
			os.Exit(1)
		}
		if err := runInspectPipeline(cfg, os.Args[2]); err != nil {
			logError("Inspection failed: %v", err)
			os.Exit(1)
		}
	case "restore":
		if len(os.Args) < 3 {
			logError("ERROR: No backup filename/S3 URL provided!")
			fmt.Println("Usage: backup-app restore <filename.sql.gz | s3://bucket/key>")
			os.Exit(1)
		}
		if err := runRestorePipeline(cfg, os.Args[2]); err != nil {
			logError("Restore failed: %v", err)
			os.Exit(1)
		}
	default:
		fmt.Printf("Unknown mode: %s\n", mode)
		os.Exit(1)
	}
}

// --- MODES OF OPERATION ---

func initScheduler(cfg Config) {
	if cfg.Schedule == "" {
		fmt.Println("No BACKUP_SCHEDULE set.")
		// Create a channel that listens for OS shutdown signals
		sigChan := make(chan os.Signal, 1)
		signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

		<-sigChan // This blocks cleanly forever until Docker stops the container
		fmt.Println("Shutting down gracefully...")
		return
	}

	fmt.Println("Initializing backup scheduler...")
	c := cron.New(cron.WithSeconds()) // Supports standard cron or standard with seconds

	_, err := c.AddFunc(cfg.Schedule, func() {
		logInfo("Scheduled job triggered.")
		if err := runBackupPipeline(cfg); err != nil {
			logError("Scheduled backup pipeline failed: %v", err)
		}
	})
	if err != nil {
		logError("Failed to parse cron schedule '%s': %v", cfg.Schedule, err)
		os.Exit(1)
	}

	fmt.Printf("Cron scheduled with pattern: %s\n", cfg.Schedule)
	fmt.Println("Starting native Go cron engine...")
	c.Run()
}

func runBackupPipeline(cfg Config) error {
	timestamp := time.Now().Format("2006-01-02_15-04-05")
	fileName := fmt.Sprintf("%s-%s.sql.gz", cfg.DB, timestamp)
	fullPath := filepath.Join(cfg.BackupDir, fileName)

	logInfo("Starting database backup...")
	if err := os.MkdirAll(cfg.BackupDir, os.ModePerm); err != nil {
		return fmt.Errorf("failed to create backup dir: %w", err)
	}

	// 1. Generate Local Backup (Stream pg_dump through Go's gzip engine)
	outFile, err := os.Create(fullPath)
	if err != nil {
		return fmt.Errorf("failed to create backup file: %w", err)
	}

	gzipWriter := gzip.NewWriter(outFile)

	args := []string{
		"-h", cfg.Host, "-p", cfg.Port, "-U", cfg.User, cfg.DB,
		"--exclude-table-data=(media_storage|media_refs|tiles_cache)",
	}
	cmd := exec.Command("pg_dump", args...)
	cmd.Env = append(os.Environ(), "PGPASSWORD="+cfg.Password)
	cmd.Stdout = gzipWriter

	if err := cmd.Run(); err != nil {
		gzipWriter.Close()
		outFile.Close()
		os.Remove(fullPath)
		return fmt.Errorf("pg_dump execution failed: %w", err)
	}
	gzipWriter.Close()
	outFile.Close()
	logInfo("Success: Created local backup at %s", fullPath)

	// 2. INTEGRITY TEST
	logInfo("Launching backup integrity test...")
	testDB := fmt.Sprintf("test_verify_%s", cfg.DB)

	_ = runPsqlCmd(cfg, "postgres", fmt.Sprintf("DROP DATABASE IF EXISTS %s;", testDB))
	if err := runPsqlCmd(cfg, "postgres", fmt.Sprintf("CREATE DATABASE %s;", testDB)); err != nil {
		return fmt.Errorf("integrity test setup failed: %w", err)
	}

	// Stream gunzip restore to test DB
	if err := streamRestoreFile(cfg, fullPath, testDB); err != nil {
		logError("CRITICAL WARNING: Backup file was created but FAILED the restoration test! File is corrupted.")
		_ = runPsqlCmd(cfg, "postgres", fmt.Sprintf("DROP DATABASE IF EXISTS %s;", testDB))
		return fmt.Errorf("integrity verification failed: %w", err)
	}

	logInfo("PASSED: Backup file successfully passed restoration test syntax verification.")
	_ = runPsqlCmd(cfg, "postgres", fmt.Sprintf("DROP DATABASE %s;", testDB))

	// 3. S3 UPLOAD (if AWS config is present)
	if cfg.AWSAccessKey != "" {
		logInfo("S3 configuration detected. Uploading to cloud storage...")
		if err := uploadToS3(cfg, fullPath, fileName); err != nil {
			return fmt.Errorf("S3 synchronization failed: %w", err)
		}
		logInfo("Cloud sync success: Uploaded to s3://%s/", cfg.S3Bucket)
	}

	// 4. RETENTION MAINTENANCE
	logInfo("Evaluating retention policy (Keeping the latest %d backups)...", cfg.KeepCount)
	if err := maintainRetention(cfg); err != nil {
		logError("Retention maintenance error: %v", err)
	}

	return nil
}

func runInspectPipeline(cfg Config, target string) error {
	localPath := target
	var err error

	if strings.HasPrefix(target, "s3://") {
		logInfo("Target is on S3. Fetching metadata header...")
		bucket, key := parseS3URI(target)
		return inspectS3File(cfg, bucket, key)
	}

	if !filepath.IsAbs(localPath) {
		localPath = filepath.Join(cfg.BackupDir, target)
	}

	file, err := os.Open(localPath)
	if err != nil {
		return fmt.Errorf("unable to open backup file: %w", err)
	}
	defer file.Close()

	stat, err := file.Stat()
	if err != nil {
		return err
	}

	gzReader, err := gzip.NewReader(file)
	if err != nil {
		return fmt.Errorf("invalid gzip compression format: %w", err)
	}
	defer gzReader.Close()

	fmt.Printf("\n--- BACKUP INSPECTION REPORT ---\n")
	fmt.Printf("File Source:      %s\n", filepath.Base(localPath))
	fmt.Printf("Compressed Size:  %.2f MB\n", float64(stat.Size())/(1024*1024))
	if !gzReader.ModTime.IsZero() {
		fmt.Printf("Gzip Timestamp:   %s\n", gzReader.ModTime.Format(time.RFC1123))
	}

	// Read first 2KB to look for pg_dump header details
	buffer := make([]byte, 2048)
	n, _ := io.ReadFull(gzReader, buffer)
	if n > 0 {
		content := string(buffer[:n])
		fmt.Println("\nHeader Snippet (First 5 comment lines):")
		lines := strings.Split(content, "\n")
		count := 0
		for _, line := range lines {
			if strings.HasPrefix(line, "--") {
				fmt.Printf("  %s\n", line)
				count++
				if count >= 5 {
					break
				}
			}
		}
	}
	fmt.Printf("--------------------------------\n\n")
	return nil
}

func runRestorePipeline(cfg Config, target string) error {
	localPath := target

	// Handle S3 targets seamlessly
	if strings.HasPrefix(target, "s3://") {
		logInfo("S3 URI detected (%s). Downloading archive to disk buffer...", target)
		bucket, key := parseS3URI(target)

		tmpFile, err := os.CreateTemp(cfg.BackupDir, "s3-restore-*.sql.gz")
		if err != nil {
			return fmt.Errorf("failed to create secure disk temporary buffer: %w", err)
		}
		defer os.Remove(tmpFile.Name()) // Auto cleanup
		defer tmpFile.Close()

		if err := downloadFromS3(cfg, bucket, key, tmpFile); err != nil {
			return fmt.Errorf("S3 file extraction streaming failed: %w", err)
		}
		localPath = tmpFile.Name()
		logInfo("Download finalized. Moving to verification sequence...")
	} else if !filepath.IsAbs(localPath) {
		localPath = filepath.Join(cfg.BackupDir, target)
	}

	if _, err := os.Stat(localPath); os.IsNotExist(err) {
		return fmt.Errorf("file not found at target resolution scope: %s", localPath)
	}

	logInfo("STEP 1: Verifying integrity of backup file inside sandbox database...")
	sandboxDB := fmt.Sprintf("restore_sandbox_%s", cfg.DB)

	_ = runPsqlCmd(cfg, "postgres", fmt.Sprintf("DROP DATABASE IF EXISTS %s;", sandboxDB))
	if err := runPsqlCmd(cfg, "postgres", fmt.Sprintf("CREATE DATABASE %s;", sandboxDB)); err != nil {
		return fmt.Errorf("sandbox creation failed: %w", err)
	}

	if err := streamRestoreFile(cfg, localPath, sandboxDB); err != nil {
		logError("CRITICAL FAILURE: The backup file failed verification tests! Main database will not be touched.")
		_ = runPsqlCmd(cfg, "postgres", fmt.Sprintf("DROP DATABASE IF EXISTS %s;", sandboxDB))
		return fmt.Errorf("sandbox restore test failed: %w", err)
	}
	logInfo("PASSED: Verification successful. Sandbox database populated cleanly.")
	_ = runPsqlCmd(cfg, "postgres", fmt.Sprintf("DROP DATABASE %s;", sandboxDB))

	logInfo("STEP 2: Proceeding with Live Database Restoration...")
	logInfo("WARNING: Dropping active public schema on target database: %s...", cfg.DB)
	if err := runPsqlCmd(cfg, cfg.DB, "DROP SCHEMA public CASCADE; CREATE SCHEMA public;"); err != nil {
		return fmt.Errorf("failed to flush target schema: %w", err)
	}

	logInfo("STEP 3: Streaming verified records into production engine...")
	if err := streamRestoreFile(cfg, localPath, cfg.DB); err != nil {
		return fmt.Errorf("critical engine restoration stream failed: %w", err)
	}

	logInfo("SUCCESS: Production database successfully restored to status: %s", filepath.Base(target))
	return nil
}

// --- HELPERS & UTILITIES ---

func runPsqlCmd(cfg Config, dbname, sqlCommand string) error {
	args := []string{"-h", cfg.Host, "-p", cfg.Port, "-U", cfg.User, "-d", dbname, "-c", sqlCommand}
	cmd := exec.Command("psql", args...)
	cmd.Env = append(os.Environ(), "PGPASSWORD="+cfg.Password)
	return cmd.Run()
}

func streamRestoreFile(cfg Config, gzippedFilePath, targetDB string) error {
	file, err := os.Open(gzippedFilePath)
	if err != nil {
		return err
	}
	defer file.Close()

	gzReader, err := gzip.NewReader(file)
	if err != nil {
		return err
	}
	defer gzReader.Close()

	args := []string{"-h", cfg.Host, "-p", cfg.Port, "-U", cfg.User, "-d", targetDB}
	cmd := exec.Command("psql", args...)
	cmd.Env = append(os.Environ(), "PGPASSWORD="+cfg.Password)

	stdinPipe, err := cmd.StdinPipe()
	if err != nil {
		return err
	}

	if err := cmd.Start(); err != nil {
		return err
	}

	if _, err := io.Copy(stdinPipe, gzReader); err != nil {
		stdinPipe.Close()
		return err
	}
	stdinPipe.Close()

	return cmd.Wait()
}

func getS3Client(cfg Config) (*s3.Client, error) {
	awsCfg, err := config.LoadDefaultConfig(context.TODO(),
		config.WithRegion(cfg.AWSRegion),
		config.WithCredentialsProvider(credentials.NewStaticCredentialsProvider(cfg.AWSAccessKey, cfg.AWSSecretKey, "")),
	)
	if err != nil {
		return nil, err
	}

	client := s3.NewFromConfig(awsCfg, func(o *s3.Options) {
		if cfg.S3Endpoint != "" {
			o.BaseEndpoint = aws.String(cfg.S3Endpoint)
			o.UsePathStyle = true
		}
	})
	return client, nil
}

func uploadToS3(cfg Config, localPath, s3Key string) error {
	file, err := os.Open(localPath)
	if err != nil {
		return err
	}
	defer file.Close()

	// 1. Load the default credentials and region configuration
	awsCfg, err := config.LoadDefaultConfig(context.TODO(),
		config.WithRegion(cfg.AWSRegion),
		config.WithCredentialsProvider(credentials.NewStaticCredentialsProvider(cfg.AWSAccessKey, cfg.AWSSecretKey, "")),
	)
	if err != nil {
		return err
	}

	// 2. Initialize the S3 Client with the modern BaseEndpoint design pattern
	client := s3.NewFromConfig(awsCfg, func(o *s3.Options) {
		if cfg.S3Endpoint != "" {
			// This string replaces the old EndpointResolver logic perfectly
			o.BaseEndpoint = aws.String(cfg.S3Endpoint)

			// If you use MinIO or DigitalOcean, you usually want virtual-host style path addressing disabled
			o.UsePathStyle = true
		}
	})

	// 3. Dispatch the payload to your cloud storage
	_, err = client.PutObject(context.TODO(), &s3.PutObjectInput{
		Bucket: aws.String(cfg.S3Bucket),
		Key:    aws.String(s3Key),
		Body:   file,
	})
	return err
}
func downloadFromS3(cfg Config, bucket, key string, targetFile *os.File) error {
	client, err := getS3Client(cfg)
	if err != nil {
		return err
	}

	result, err := client.GetObject(context.TODO(), &s3.GetObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(key),
	})
	if err != nil {
		return err
	}
	defer result.Body.Close()

	_, err = io.Copy(targetFile, result.Body)
	return err
}
func inspectS3File(cfg Config, bucket, key string) error {
	client, err := getS3Client(cfg)
	if err != nil {
		return err
	}

	result, err := client.GetObject(context.TODO(), &s3.GetObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(key),
	})
	if err != nil {
		return fmt.Errorf("failed to fetch target from S3: %w", err)
	}
	defer result.Body.Close()

	gzReader, err := gzip.NewReader(result.Body)
	if err != nil {
		return fmt.Errorf("downloaded target is not a valid gzip configuration file: %w", err)
	}
	defer gzReader.Close()

	fmt.Printf("\n--- S3 REMOTE BACKUP INSPECTION REPORT ---\n")
	fmt.Printf("Bucket target:    s3://%s/%s\n", bucket, key)
	if result.ContentLength != nil {
		fmt.Printf("Compressed Size:  %.2f MB\n", float64(*result.ContentLength)/(1024*1024))
	}
	if result.LastModified != nil {
		fmt.Printf("S3 Modification:  %s\n", result.LastModified.Format(time.RFC1123))
	}

	buffer := make([]byte, 2048)
	n, _ := io.ReadFull(gzReader, buffer)
	if n > 0 {
		fmt.Println("\nHeader Snippet (First 5 comment lines):")
		lines := strings.Split(string(buffer[:n]), "\n")
		count := 0
		for _, line := range lines {
			if strings.HasPrefix(line, "--") {
				fmt.Printf("  %s\n", line)
				count++
				if count >= 5 {
					break
				}
			}
		}
	}
	fmt.Printf("------------------------------------------\n\n")
	return nil
}
func maintainRetention(cfg Config) error {
	pattern := filepath.Join(cfg.BackupDir, fmt.Sprintf("%s-*.sql.gz", cfg.DB))
	files, err := filepath.Glob(pattern)
	if err != nil {
		return err
	}

	logInfo("Found %d total backup files locally.", len(files))
	if len(files) <= cfg.KeepCount {
		logInfo("No files require rotation.")
		return nil
	}

	// Sort files alphabetically (since naming includes timestamps, alpha sort maps cleanly to chronological sort)
	sort.Strings(files)

	purgeCount := len(files) - cfg.KeepCount
	logInfo("Threshold exceeded. Removing oldest %d file(s)...", purgeCount)

	for i := 0; i < purgeCount; i++ {
		logInfo("Deleted expired file: %s", filepath.Base(files[i]))
		if err := os.Remove(files[i]); err != nil {
			logError("Failed to delete %s: %v", files[i], err)
		}
	}
	logInfo("Retention cleanup process finished.")
	return nil
}
func parseS3URI(uri string) (bucket, key string) {
	cleanStr := strings.TrimPrefix(uri, "s3://")
	parts := strings.SplitN(cleanStr, "/", 2)
	if len(parts) == 2 {
		return parts[0], parts[1]
	}
	return parts[0], ""
}

func getEnv(key, fallback string) string {
	if value, exists := os.LookupEnv(key); exists {
		return value
	}
	return fallback
}

func logInfo(format string, v ...interface{}) {
	fmt.Printf("[%s] %s\n", time.Now().Format("2006-01-02 15:04:05"), fmt.Sprintf(format, v...))
}

func logError(format string, v ...interface{}) {
	fmt.Fprintf(os.Stderr, "[%s] ERROR: %s\n", time.Now().Format("2006-01-02 15:04:05"), fmt.Sprintf(format, v...))
}

func printHelp() {
	fmt.Println(`===================================================================
Database Backup Manager (Go Edition)
===================================================================
Usage: 
  backup-app <command> [arguments]

Commands:
  cron-init          Starts the persistent background cron engine using the
                     cron pattern defined in $BACKUP_SCHEDULE. If no schedule
                     is defined, the process waits cleanly for termination.
  
  run                Triggers a single manual backup sequence immediately:
                     Dumps DB -> Compresses -> Verifies -> Uploads to S3 -> Rotates.

  inspect <target>   Peeks into a local archive file or remote S3 object URI
                     to extract GZIP headers, modification timestamps, and the
                     first few lines of pg_dump comments without a full download.
                     Examples:
                       backup-app inspect my-backup.sql.gz
                       backup-app inspect s3://my-bucket/my-backup.sql.gz
  
  restore <target>   Performs a safe, multi-stage restore of a local file or 
                     S3 object URI. Streams to a sandbox DB for verification 
                     before wiping and updating the production database schema.
                     Examples:
                       backup-app restore my-backup.sql.gz
                       backup-app restore s3://my-bucket/my-backup.sql.gz

  help, -h, --help   Display this help screen.

Required Environment Variables:
  POSTGRES_HOST         Database server hostname/IP
  POSTGRES_USER         Database administrative user
  POSTGRES_PASSWORD     Database password
  POSTGRES_DB           Target database name to back up/restore

Optional Environment Variables:
  POSTGRES_PORT         Database port (Default: 5432)
  BACKUP_DIR            Local storage path for backups (Default: /backups)
  BACKUP_KEEP_COUNT     Number of local archives to retain (Default: 7)
  BACKUP_SCHEDULE       Cron interval string (e.g., "0 2 * * *")
  
S3 Cloud Variables (Required for s3:// commands):
  AWS_ACCESS_KEY_ID     Your IAM Access Key ID
  AWS_SECRET_ACCESS_KEY Your IAM Secret Access Key
  S3_BUCKET             Target S3 Bucket name (used during 'run')
  AWS_DEFAULT_REGION    Target AWS region (Default: us-east-1)
  S3_ENDPOINT           Custom URL for S3-compatible APIs (MinIO, DigitalOcean)
===================================================================`)
}