// auth_service_account — Google Service Account JWT auth primitive
//
// 讀取 auth_recipe:{service} + 解密 service_account_json + 組 JWT + RS256 簽章(透過 host)
// + token exchange → access_token + 展開 {{runtime.access_token}}。
//
// Host imports:
//   - u6u.kv_get              — 讀 RECIPES + CREDENTIALS_KV
//   - u6u.crypto_decrypt      — AES-GCM 解密 service account JSON
//   - u6u.crypto_sign_rs256   — RSASSA-PKCS1-v1_5 + SHA-256 (PKCS8 private key)
//   - u6u.http_request        — POST token exchange endpoint
//
//go:build tinygo

package main

import (
	"encoding/base64"
	"encoding/json"
	"io"
	"net/url"
	"os"
	"strings"
	"time"
	"unsafe"
)

// ── host function 宣告 ───────────────────────────────────────────────────────

//go:wasmimport u6u kv_get
func hostKvGet(
	keyPtr uintptr, keyLen uint32,
	outPtr uintptr, outLenPtr uintptr,
) uint32

//go:wasmimport u6u crypto_decrypt
func hostCryptoDecrypt(
	encPtr uintptr, encLen uint32,
	ivPtr uintptr, ivLen uint32,
	outPtr uintptr, outLenPtr uintptr,
) uint32

//go:wasmimport u6u crypto_sign_rs256
func hostCryptoSignRS256(
	dataPtr uintptr, dataLen uint32,
	pkcs8Ptr uintptr, pkcs8Len uint32,
	outPtr uintptr, outLenPtr uintptr,
) uint32

//go:wasmimport u6u http_request
func hostHttpRequest(
	urlPtr uintptr, urlLen uint32,
	methodPtr uintptr, methodLen uint32,
	headersPtr uintptr, headersLen uint32,
	bodyPtr uintptr, bodyLen uint32,
	outPtr uintptr, outLenPtr uintptr,
) uint32

// ── 型別 ─────────────────────────────────────────────────────────────────────

type Input struct {
	Action  string          `json:"action"`
	APIKey  string          `json:"api_key"`
	Service string          `json:"service"`
	Request json.RawMessage `json:"request,omitempty"`
}

type SecretRequirement struct {
	Key      string `json:"key"`
	Label    string `json:"label"`
	Optional bool   `json:"optional,omitempty"`
}

type AuthInjectSpec struct {
	Header map[string]string `json:"header,omitempty"`
	Query  map[string]string `json:"query,omitempty"`
	Body   map[string]string `json:"body,omitempty"`
}

type TokenExchange struct {
	Endpoint string   `json:"endpoint"`
	Scopes   []string `json:"scopes"`
}

type AuthRecipe struct {
	Kind               string              `json:"kind"`
	Service            string              `json:"service"`
	Primitive          string              `json:"primitive"`
	ServiceAccountKind string              `json:"service_account_kind,omitempty"`
	TokenExchange      *TokenExchange      `json:"token_exchange,omitempty"`
	RequiredSecrets    []SecretRequirement `json:"required_secrets"`
	Inject             AuthInjectSpec      `json:"inject"`
}

type EncryptedRecord struct {
	Encrypted string `json:"encrypted"`
	IV        string `json:"iv"`
}

type ServiceAccountJSON struct {
	ClientEmail string `json:"client_email"`
	PrivateKey  string `json:"private_key"`
}

type JWTHeader struct {
	Alg string `json:"alg"`
	Typ string `json:"typ"`
}

type JWTPayload struct {
	Iss   string `json:"iss"`
	Sub   string `json:"sub"`
	Aud   string `json:"aud"`
	Scope string `json:"scope"`
	Iat   int64  `json:"iat"`
	Exp   int64  `json:"exp"`
}

// ── main ─────────────────────────────────────────────────────────────────────

func main() {
	raw, err := io.ReadAll(os.Stdin)
	if err != nil {
		writeError("failed to read stdin: " + err.Error())
		return
	}

	var input Input
	if err := json.Unmarshal(raw, &input); err != nil {
		writeError("invalid input JSON: " + err.Error())
		return
	}

	if input.APIKey == "" {
		writeError("api_key 必填")
		return
	}
	if input.Service == "" {
		writeError("service 必填")
		return
	}
	if input.Action != "" && input.Action != "authenticate" {
		writeError("auth_service_account 僅支援 action=authenticate")
		return
	}

	// 1. 讀 auth recipe
	recipeJSON, status := kvGet("auth_recipe:" + input.Service)
	if status == 2 {
		writeError("找不到 auth recipe: " + input.Service)
		return
	}
	if status != 0 {
		writeError("kv_get 失敗(auth_recipe)")
		return
	}

	var recipe AuthRecipe
	if err := json.Unmarshal([]byte(recipeJSON), &recipe); err != nil {
		writeError("auth recipe JSON 解析失敗: " + err.Error())
		return
	}
	if recipe.Primitive != "service_account" {
		writeError("auth recipe " + input.Service + " 的 primitive 不是 service_account (是 " + recipe.Primitive + ")")
		return
	}
	if recipe.ServiceAccountKind != "google_jwt" {
		writeError("auth recipe " + input.Service + " 的 service_account_kind 必須是 google_jwt,實際: " + recipe.ServiceAccountKind)
		return
	}
	if recipe.TokenExchange == nil || recipe.TokenExchange.Endpoint == "" {
		writeError("auth recipe " + input.Service + " 缺少 token_exchange.endpoint")
		return
	}
	if len(recipe.RequiredSecrets) == 0 {
		writeError("auth recipe " + input.Service + " 缺少 required_secrets[0](SA JSON)")
		return
	}

	// 2. 解密 service account JSON (慣例:required_secrets[0] 是 SA JSON)
	saReq := recipe.RequiredSecrets[0]
	kvKey := input.APIKey + ":cred:" + saReq.Key
	encJSON, s := kvGet(kvKey)
	if s == 2 {
		writeError("缺少 credential: " + saReq.Key + " (" + saReq.Label + ")。修復: 編輯 credentials.yaml 後執行 acr creds push")
		return
	}
	if s != 0 {
		writeError("kv_get 失敗(credential " + saReq.Key + ")")
		return
	}

	var rec EncryptedRecord
	if err := json.Unmarshal([]byte(encJSON), &rec); err != nil {
		writeError("credential " + saReq.Key + " 格式錯誤: " + err.Error())
		return
	}
	saJSONStr, ok := cryptoDecrypt(rec.Encrypted, rec.IV)
	if !ok {
		writeError("credential " + saReq.Key + " 解密失敗")
		return
	}

	// 3. 解析 service account JSON
	var sa ServiceAccountJSON
	if err := json.Unmarshal([]byte(saJSONStr), &sa); err != nil {
		writeError("service account JSON 格式錯誤: " + err.Error())
		return
	}
	if sa.ClientEmail == "" || sa.PrivateKey == "" {
		writeError("service account JSON 缺少 client_email 或 private_key")
		return
	}

	// 4. PEM → PKCS8 bytes (去 header/footer + base64 decode)
	pkcs8, err := pemToPkcs8(sa.PrivateKey)
	if err != nil {
		writeError("解析 service account private key 失敗: " + err.Error())
		return
	}

	// 5. 組 JWT header + payload (base64url-encoded)
	now := time.Now().Unix()
	header := JWTHeader{Alg: "RS256", Typ: "JWT"}
	payload := JWTPayload{
		Iss:   sa.ClientEmail,
		Sub:   sa.ClientEmail,
		Aud:   recipe.TokenExchange.Endpoint,
		Scope: strings.Join(recipe.TokenExchange.Scopes, " "),
		Iat:   now,
		Exp:   now + 3600,
	}
	headerBytes, _ := json.Marshal(header)
	payloadBytes, _ := json.Marshal(payload)

	signingInput := base64.RawURLEncoding.EncodeToString(headerBytes) + "." +
		base64.RawURLEncoding.EncodeToString(payloadBytes)

	// 6. 呼叫 host 簽章 (RSASSA-PKCS1-v1_5 + SHA-256)
	signature, ok := cryptoSignRS256([]byte(signingInput), pkcs8)
	if !ok {
		writeError("JWT 簽章失敗(host function crypto_sign_rs256 回傳錯誤)")
		return
	}

	jwt := signingInput + "." + base64.RawURLEncoding.EncodeToString(signature)

	// 7. token exchange:POST form-urlencoded 到 token_exchange.endpoint
	form := url.Values{}
	form.Set("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer")
	form.Set("assertion", jwt)
	formBody := form.Encode()

	headersJSON := `{"Content-Type":"application/x-www-form-urlencoded"}`

	respStr, ok := httpRequest(recipe.TokenExchange.Endpoint, "POST", headersJSON, formBody)
	if !ok {
		writeError("token exchange HTTP 失敗")
		return
	}

	var tokenResp struct {
		AccessToken string `json:"access_token"`
		Error       string `json:"error"`
		ErrorDesc   string `json:"error_description"`
	}
	if err := json.Unmarshal([]byte(respStr), &tokenResp); err != nil {
		writeError("token exchange 回應解析失敗: " + err.Error() + " (raw: " + respStr + ")")
		return
	}
	if tokenResp.AccessToken == "" {
		errMsg := tokenResp.Error
		if tokenResp.ErrorDesc != "" {
			errMsg += ": " + tokenResp.ErrorDesc
		}
		if errMsg == "" {
			errMsg = "access_token 為空 (raw: " + respStr + ")"
		}
		writeError("token exchange 失敗: " + errMsg)
		return
	}

	// 8. 展開模板 (service_account 不用 secret.*,只用 runtime.access_token)
	secrets := map[string]string{}
	runtime := map[string]string{"access_token": tokenResp.AccessToken}
	authHeaders := interpolateRecord(recipe.Inject.Header, secrets, runtime)
	authQuery := interpolateRecord(recipe.Inject.Query, secrets, runtime)
	authBody := interpolateRecord(recipe.Inject.Body, secrets, runtime)

	out, _ := json.Marshal(map[string]interface{}{
		"success":      true,
		"auth_headers": authHeaders,
		"auth_query":   authQuery,
		"auth_body":    authBody,
		"runtime":      runtime,
	})
	os.Stdout.Write(out)
}

// ── helpers ──────────────────────────────────────────────────────────────────

func writeError(msg string) {
	out, _ := json.Marshal(map[string]interface{}{
		"success":      false,
		"error":        msg,
		"auth_headers": map[string]string{},
		"auth_query":   map[string]string{},
		"auth_body":    map[string]string{},
	})
	os.Stdout.Write(out)
}

// pemToPkcs8 從 PEM 取出 base64 body 再 decode 成 bytes。
// 支援 "BEGIN PRIVATE KEY" / "BEGIN RSA PRIVATE KEY"(SA JSON 幾乎都是前者)。
func pemToPkcs8(pem string) ([]byte, error) {
	// 移除所有 BEGIN/END 行與空白
	lines := strings.Split(pem, "\n")
	var b strings.Builder
	for _, line := range lines {
		l := strings.TrimSpace(line)
		if l == "" {
			continue
		}
		if strings.HasPrefix(l, "-----BEGIN") || strings.HasPrefix(l, "-----END") {
			continue
		}
		b.WriteString(l)
	}
	cleaned := strings.ReplaceAll(b.String(), "\\n", "") // 防呆:JSON-escaped newline
	return base64.StdEncoding.DecodeString(cleaned)
}

// kvGet 呼叫 host function,回傳 (value, status)。status: 0=成功 1=錯誤 2=找不到
func kvGet(key string) (string, uint32) {
	keyBytes := []byte(key)
	outBuf := make([]byte, 65536)
	var outLen uint32

	status := hostKvGet(
		uintptr(unsafe.Pointer(&keyBytes[0])), uint32(len(keyBytes)),
		uintptr(unsafe.Pointer(&outBuf[0])), uintptr(unsafe.Pointer(&outLen)),
	)
	if status != 0 {
		return "", status
	}
	return string(outBuf[:outLen]), 0
}

func cryptoDecrypt(encB64, ivB64 string) (string, bool) {
	encBytes := []byte(encB64)
	ivBytes := []byte(ivB64)
	outBuf := make([]byte, 65536)
	var outLen uint32

	if len(encBytes) == 0 || len(ivBytes) == 0 {
		return "", false
	}

	status := hostCryptoDecrypt(
		uintptr(unsafe.Pointer(&encBytes[0])), uint32(len(encBytes)),
		uintptr(unsafe.Pointer(&ivBytes[0])), uint32(len(ivBytes)),
		uintptr(unsafe.Pointer(&outBuf[0])), uintptr(unsafe.Pointer(&outLen)),
	)
	if status != 0 {
		return "", false
	}
	return string(outBuf[:outLen]), true
}

// cryptoSignRS256 呼叫 host,回傳簽章 bytes
func cryptoSignRS256(data, pkcs8 []byte) ([]byte, bool) {
	if len(data) == 0 || len(pkcs8) == 0 {
		return nil, false
	}
	outBuf := make([]byte, 1024) // RSA-2048 簽章 = 256 bytes,1KB 綽綽有餘
	var outLen uint32

	status := hostCryptoSignRS256(
		uintptr(unsafe.Pointer(&data[0])), uint32(len(data)),
		uintptr(unsafe.Pointer(&pkcs8[0])), uint32(len(pkcs8)),
		uintptr(unsafe.Pointer(&outBuf[0])), uintptr(unsafe.Pointer(&outLen)),
	)
	if status != 0 {
		return nil, false
	}
	return outBuf[:outLen], true
}

// httpRequest 呼叫 host,回傳 response body 字串(host 側把 status + body 串好)
func httpRequest(url, method, headersJSON, body string) (string, bool) {
	urlBytes := []byte(url)
	methodBytes := []byte(method)
	headersBytes := []byte(headersJSON)
	bodyBytes := []byte(body)

	if len(urlBytes) == 0 {
		return "", false
	}

	outBuf := make([]byte, 65536)
	var outLen uint32

	// bodyBytes 可能為空(GET),host function 允許 len=0
	var bodyPtr uintptr
	if len(bodyBytes) > 0 {
		bodyPtr = uintptr(unsafe.Pointer(&bodyBytes[0]))
	}
	var headersPtr uintptr
	if len(headersBytes) > 0 {
		headersPtr = uintptr(unsafe.Pointer(&headersBytes[0]))
	}

	status := hostHttpRequest(
		uintptr(unsafe.Pointer(&urlBytes[0])), uint32(len(urlBytes)),
		uintptr(unsafe.Pointer(&methodBytes[0])), uint32(len(methodBytes)),
		headersPtr, uint32(len(headersBytes)),
		bodyPtr, uint32(len(bodyBytes)),
		uintptr(unsafe.Pointer(&outBuf[0])), uintptr(unsafe.Pointer(&outLen)),
	)
	if status != 0 {
		return "", false
	}
	return string(outBuf[:outLen]), true
}

// interpolateTemplate 展開 {{secret.X}} 與 {{runtime.X}}。未知 key 展開為空字串。
// 其他 namespace 的 {{...}} 原樣保留。
func interpolateTemplate(template string, secrets, runtime map[string]string) string {
	var b strings.Builder
	b.Grow(len(template))
	i := 0
	for i < len(template) {
		start := strings.Index(template[i:], "{{")
		if start < 0 {
			b.WriteString(template[i:])
			break
		}
		b.WriteString(template[i : i+start])
		openIdx := i + start
		closeRel := strings.Index(template[openIdx+2:], "}}")
		if closeRel < 0 {
			b.WriteString(template[openIdx:])
			break
		}
		inner := template[openIdx+2 : openIdx+2+closeRel]
		advance := openIdx + 2 + closeRel + 2

		switch {
		case strings.HasPrefix(inner, "secret."):
			key := inner[len("secret."):]
			b.WriteString(secrets[key])
		case strings.HasPrefix(inner, "runtime."):
			key := inner[len("runtime."):]
			b.WriteString(runtime[key])
		default:
			b.WriteString(template[openIdx:advance])
		}
		i = advance
	}
	return b.String()
}

func interpolateRecord(
	record map[string]string,
	secrets, runtime map[string]string,
) map[string]string {
	if record == nil {
		return map[string]string{}
	}
	result := make(map[string]string, len(record))
	for k, v := range record {
		result[k] = interpolateTemplate(v, secrets, runtime)
	}
	return result
}