package main

import (
    "context"
    "fmt"

    "github.com/aws/aws-sdk-go-v2/aws"
    "github.com/aws/aws-sdk-go-v2/config"
    "github.com/aws/aws-sdk-go-v2/service/ec2"
)

// SnapshotReport holds the volume ID and the count of its snapshots.
type SnapshotReport struct {
    VolumeID       string
    SnapshotCount  int
}

// GenerateSnapshotReport retrieves all snapshots and counts how many snapshots each volume has.
func GenerateSnapshotReport(ctx context.Context) ([]SnapshotReport, error) {
    cfg, err := config.LoadDefaultConfig(ctx)
    if err != nil {
        return nil, fmt.Errorf("failed to load AWS SDK config: %w", err)
    }

    ec2Client := ec2.NewFromConfig(cfg)
    paginator := ec2.NewDescribeSnapshotsPaginator(ec2Client, &ec2.DescribeSnapshotsInput{
        OwnerIds: []string{"self"},
    })

    volumeSnapshotCounts := make(map[string]int)
    fmt.Println("Fetching snapshots and calculating volume snapshot counts...")
    for paginator.HasMorePages() {
        out, err := paginator.NextPage(ctx)
        if err != nil {
            return nil, fmt.Errorf("failed to fetch snapshots: %w", err)
        }

        for _, snapshot := range out.Snapshots {
            volumeID := *snapshot.VolumeId
            volumeSnapshotCounts[volumeID]++
        }
    }

    var reports []SnapshotReport
    for volumeID, count := range volumeSnapshotCounts {
        reports = append(reports, SnapshotReport{
            VolumeID:      volumeID,
            SnapshotCount: count,
        })
    }

    return reports, nil
}

func main() {
    ctx := context.Background()
    reports, err := GenerateSnapshotReport(ctx)
    if err != nil {
        fmt.Printf("Error: %s\n", err)
        return
    }

    fmt.Println("Snapshot Report:")
    for _, report := range reports {
        fmt.Printf("Volume ID: %s, Snapshot Count: %d\n", report.VolumeID, report.SnapshotCount)
    }
}