package main

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/credentials"
	"github.com/aws/aws-sdk-go-v2/service/sts"
)

// GCE metadata ID token (audience = the AWS role ARN, format=full so azp/sub are present) -> AssumeRoleWithWebIdentity
// (Terraform design 6.11; the trust policy accepts aud in [role ARN, SA unique id]).
func gceIDToken(ctx context.Context, audience string) (string, error) {
	u := "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity?format=full&audience=" + url.QueryEscape(audience)
	req, _ := http.NewRequestWithContext(ctx, "GET", u, nil)
	req.Header.Set("Metadata-Flavor", "Google")
	resp, err := (&http.Client{Timeout: 10 * time.Second}).Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	b, _ := io.ReadAll(resp.Body)
	if resp.StatusCode != 200 {
		return "", fmt.Errorf("metadata identity: %s: %s", resp.Status, string(b))
	}
	return string(b), nil
}

func (a *App) awsCreds(ctx context.Context) (aws.Config, error) {
	if a.cfg.AWSRoleARN == "" {
		return aws.Config{}, fmt.Errorf("AWS_ROLE_ARN not set")
	}
	tok, err := gceIDToken(ctx, a.cfg.AWSRoleARN)
	if err != nil {
		return aws.Config{}, err
	}
	anon := aws.Config{Region: a.cfg.AWSRegion, Credentials: aws.AnonymousCredentials{}}
	out, err := sts.NewFromConfig(anon).AssumeRoleWithWebIdentity(ctx, &sts.AssumeRoleWithWebIdentityInput{
		RoleArn: aws.String(a.cfg.AWSRoleARN), RoleSessionName: aws.String("harness-runner"), WebIdentityToken: aws.String(tok), DurationSeconds: aws.Int32(3600),
	})
	if err != nil {
		return aws.Config{}, fmt.Errorf("assume-role-with-web-identity: %w", err)
	}
	c := out.Credentials
	return aws.Config{Region: a.cfg.AWSRegion, Credentials: credentials.NewStaticCredentialsProvider(*c.AccessKeyId, *c.SecretAccessKey, *c.SessionToken)}, nil
}

func (a *App) awsCallerIdentity(ctx context.Context) (string, error) {
	cfg, err := a.awsCreds(ctx)
	if err != nil {
		return "", err
	}
	id, err := sts.NewFromConfig(cfg).GetCallerIdentity(ctx, &sts.GetCallerIdentityInput{})
	if err != nil {
		return "", err
	}
	return aws.ToString(id.Arn), nil
}