package controllers

import (
	"bytes"
	"net/http"
	"text/template"
	"time"

	"code.crute.us/mcrute/cloud-identity-broker/cloud/aws"

	"code.crute.us/mcrute/golib/echo/controller"
	"github.com/labstack/echo/v4"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
)

var credsAllowed = promauto.NewCounterVec(prometheus.CounterOpts{
	Namespace: "aws_access", // Legacy namespace
	Name:      "broker_cred_access_total",
	Help:      "Total number of credential accesses allowed by broker",
}, []string{"account", "region"})

type jsonCredential struct {
	AccessKeyId     *string    `json:"access_key"`
	SecretAccessKey *string    `json:"secret_key"`
	SessionToken    *string    `json:"session_token"`
	Expiration      *time.Time `json:"expiration"`
	ShortName       string     `json:"-"`
}

var (
	bashTemplate = template.Must(template.New("").Parse(
		`export AWS_CREDS_EXPIRATION="{{ .Expiration }}"
export AWS_ACCESS_KEY_ID="{{ .AccessKeyId }}"
export AWS_SECRET_ACCESS_KEY="{{ .SecretAccessKey }}"
export AWS_SESSION_TOKEN="{{ .SessionToken }}"
`))

	pslTemplate = template.Must(template.New("").Parse(
		`Set-Item -path env:AWS_CREDS_EXPIRATION -value '{{ .Expiration }}'
Set-Item -path env:AWS_ACCESS_KEY_ID -value '{{ .AccessKeyId }}'
Set-Item -path env:AWS_SECRET_ACCESS_KEY -value '{{ .SecretAccessKey }}'
Set-Item -path env:AWS_SESSION_TOKEN -value '{{ .SessionToken }}'
`))

	iniTemplate = template.Must(template.New("").Parse(
		`[profile {{ .ShortName }}]
aws_access_key_id={{ .AccessKeyId }}
aws_secret_access_key={{ .SecretAccessKey }}
aws_session_token={{ .SessionToken }}
expiration={{ .Expiration }}
`))
)

type APICredentialsHandler struct {
	*AWSAPI
}

func NewAPICredentialsHandler(a *AWSAPI) echo.HandlerFunc {
	al := &APICredentialsHandler{a}
	h := &controller.ContentTypeNegotiatingHandler{
		DefaultHandler: al.HandleJSONV1,
		Handlers: map[string]echo.HandlerFunc{
			contentTypeV1:              al.HandleJSONV1,
			contentTypeV2:              al.HandleJSONV1,
			contentTypeV2AWSBash:       al.HandleBashV2,
			contentTypeV2AWSPowershell: al.HandlePSLV2,
			contentTypeV2AWSConfig:     al.HandleINIV2,
		},
	}
	return h.Handle
}

func (h *APICredentialsHandler) getAWSCredential(c echo.Context) (*jsonCredential, error) {
	rc, err := h.GetContext(c) // Does authorization checks
	if err != nil {
		return nil, err
	}

	region := c.Param("region")
	creds, err := rc.AWS.AssumeRole(rc.Principal.Username, &region)
	if err != nil {
		if aws.IsRegionNotExist(err) {
			return nil, echo.NotFoundHandler(c)
		}
		c.Logger().Errorf("Error retrieving credentials: %w", err)
		return nil, echo.ErrInternalServerError
	}

	c.Logger().Infof(
		"Allowing '%s' to access account credential '%s'",
		rc.Principal.Username, rc.Account.Name,
	)
	credsAllowed.With(prometheus.Labels{
		"account": rc.Account.ShortName,
		"region":  region,
	}).Inc()

	c.Response().Header().Set("Expires", creds.Expiration.Add(-5*time.Minute).Format(time.RFC1123))

	return &jsonCredential{
		AccessKeyId:     creds.AccessKeyId,
		SecretAccessKey: creds.SecretAccessKey,
		SessionToken:    creds.SessionToken,
		Expiration:      creds.Expiration,
		ShortName:       rc.Account.ShortName,
	}, nil
}

func (h *APICredentialsHandler) renderTemplate(c echo.Context, t *template.Template, ct string) error {
	creds, err := h.getAWSCredential(c)
	if err != nil {
		return err
	}

	buf := &bytes.Buffer{}
	if err = t.Execute(buf, creds); err != nil {
		return echo.ErrInternalServerError
	}

	return c.Blob(http.StatusOK, ct, buf.Bytes())
}

func (h *APICredentialsHandler) HandleJSONV1(c echo.Context) error {
	creds, err := h.getAWSCredential(c)
	if err != nil {
		return err
	}

	return c.JSON(http.StatusOK, creds)
}

func (h *APICredentialsHandler) HandleBashV2(c echo.Context) error {
	return h.renderTemplate(c, bashTemplate, contentTypeV2AWSBash)
}

func (h *APICredentialsHandler) HandlePSLV2(c echo.Context) error {
	return h.renderTemplate(c, pslTemplate, contentTypeV2AWSPowershell)
}

func (h *APICredentialsHandler) HandleINIV2(c echo.Context) error {
	return h.renderTemplate(c, iniTemplate, contentTypeV2AWSConfig)
}