// Package fargate generates standalone ECS Fargate burst-seeder task bundles.
package fargate

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
)

const (
	DefaultTaskFamily = "aetherforge-burst-seeder"
	DefaultImage      = "public.ecr.aws/docker/library/nginx:alpine"
	DefaultRegion     = "us-east-1"
	DefaultCluster    = "aetherforge-burst"
)

type Options struct {
	BuildID, Campaign, ServerURL, ShardToken, PayloadSHA256 string
	PayloadSize, DataShards, ParityShards, TTLHours, TaskCount int
	ShardURLs []string
	SwarmMagnet string
	Shards [][]byte
	Region, Cluster, TaskFamily, Image string
}

type Bundle struct {
	TaskDefinitionJSON, RunTaskScript, ShardManifestJSON []byte
	ShardToken string
}

func GenerateBundle(opts Options) (*Bundle, error) {
	opts = opts.withDefaults()
	if len(opts.Shards) == 0 {
		return nil, fmt.Errorf("fargate: shards required")
	}
	manifest, err := json.MarshalIndent(map[string]interface{}{
		"token": opts.ShardToken, "shards_b64": encodeShards(opts.Shards), "ttl_hours": opts.TTLHours,
	}, "", "  ")
	if err != nil {
		return nil, err
	}
	taskDef, err := json.MarshalIndent(map[string]interface{}{
		"family": opts.TaskFamily, "networkMode": "awsvpc", "requiresCompatibilities": []string{"FARGATE"},
		"cpu": "256", "memory": "512",
		"containerDefinitions": []map[string]interface{}{{
			"name": "burst-seeder", "image": opts.Image, "essential": true,
			"environment": []map[string]string{
				{"name": "AF_SHARD_TOKEN", "value": opts.ShardToken},
				{"name": "AF_SHARD_MANIFEST_B64", "value": base64.StdEncoding.EncodeToString(manifest)},
			},
		}},
	}, "", "  ")
	if err != nil {
		return nil, err
	}
	script := fmt.Sprintf(`#!/usr/bin/env bash
set -euo pipefail
aws ecs register-task-definition --region "${AWS_REGION:-%s}" --cli-input-json file://task-definition.json
aws ecs run-task --region "${AWS_REGION:-%s}" --cluster "${ECS_CLUSTER:-%s}" --launch-type FARGATE --task-definition %s --count ${ECS_TASK_COUNT:-%d}
`, opts.Region, opts.Region, opts.Cluster, opts.TaskFamily, opts.TaskCount)
	return &Bundle{TaskDefinitionJSON: taskDef, RunTaskScript: []byte(script), ShardManifestJSON: manifest, ShardToken: opts.ShardToken}, nil
}

func (o Options) withDefaults() Options {
	if o.TaskFamily == "" {
		o.TaskFamily = DefaultTaskFamily
	}
	if o.Image == "" {
		o.Image = DefaultImage
	}
	if o.Region == "" {
		o.Region = DefaultRegion
	}
	if o.Cluster == "" {
		o.Cluster = DefaultCluster
	}
	o.TTLHours = ClampTTLHours(o.TTLHours)
	if o.TaskCount <= 0 {
		o.TaskCount = 3
	}
	return o
}

func encodeShards(shards [][]byte) []string {
	out := make([]string, len(shards))
	for i, sh := range shards {
		out[i] = base64.StdEncoding.EncodeToString(sh)
	}
	return out
}

func ClampTTLHours(h int) int {
	if h < 2 {
		return 2
	}
	if h > 4 {
		return 4
	}
	return h
}

func BurstTaskDefinition(opts Options) ([]byte, error) {
	b, err := GenerateBundle(opts)
	if err != nil {
		return nil, err
	}
	return b.TaskDefinitionJSON, nil
}