package oidcsession

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"errors"
	"net/url"
	"strings"
	"time"
)

const LoginTransactionCookieName = "easyai_gateway_oidc_login"

var loginTransactionAAD = []byte("easyai-gateway/oidc-login-transaction/v1")

type LoginTransaction struct {
	State        string    `json:"state"`
	Nonce        string    `json:"nonce"`
	PKCEVerifier string    `json:"pkceVerifier"`
	ReturnTo     string    `json:"returnTo"`
	CreatedAt    time.Time `json:"createdAt"`
}

func NewLoginTransaction(returnTo string, now time.Time) (LoginTransaction, string, error) {
	if !ValidReturnTo(returnTo) {
		return LoginTransaction{}, "", errors.New("returnTo must be a same-origin relative path")
	}
	state, err := randomBase64URL(32)
	if err != nil {
		return LoginTransaction{}, "", err
	}
	nonce, err := randomBase64URL(32)
	if err != nil {
		return LoginTransaction{}, "", err
	}
	verifier, err := randomBase64URL(32)
	if err != nil {
		return LoginTransaction{}, "", err
	}
	challengeHash := sha256.Sum256([]byte(verifier))
	return LoginTransaction{State: state, Nonce: nonce, PKCEVerifier: verifier, ReturnTo: returnTo, CreatedAt: now.UTC()},
		base64.RawURLEncoding.EncodeToString(challengeHash[:]), nil
}

func (c *Cipher) EncodeLoginTransaction(transaction LoginTransaction) (string, error) {
	payload, err := c.SealJSON(transaction, loginTransactionAAD)
	if err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(payload), nil
}

func (c *Cipher) DecodeLoginTransaction(encoded string, now time.Time) (LoginTransaction, error) {
	payload, err := base64.RawURLEncoding.DecodeString(strings.TrimSpace(encoded))
	if err != nil {
		return LoginTransaction{}, errors.New("OIDC login transaction is invalid")
	}
	var transaction LoginTransaction
	if err := c.OpenJSON(payload, loginTransactionAAD, &transaction); err != nil {
		return LoginTransaction{}, err
	}
	if transaction.State == "" || transaction.Nonce == "" || transaction.PKCEVerifier == "" || !ValidReturnTo(transaction.ReturnTo) ||
		transaction.CreatedAt.IsZero() || now.Before(transaction.CreatedAt.Add(-time.Minute)) || !now.Before(transaction.CreatedAt.Add(10*time.Minute)) {
		return LoginTransaction{}, errors.New("OIDC login transaction has expired or is invalid")
	}
	return transaction, nil
}

func ValidReturnTo(value string) bool {
	value = strings.TrimSpace(value)
	if value == "" || !strings.HasPrefix(value, "/") || strings.HasPrefix(value, "//") || strings.Contains(value, "\\") {
		return false
	}
	parsed, err := url.Parse(value)
	return err == nil && !parsed.IsAbs() && parsed.Host == ""
}

func randomBase64URL(size int) (string, error) {
	value := make([]byte, size)
	if _, err := rand.Read(value); err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(value), nil
}