// auth_oauth2 — OAuth2 auth primitive
//
// 讀取 auth_recipe:{service}（含 token_endpoint、client_id、client_secret）
// + 解密 refresh_token + POST token endpoint 換 access_token
// + 展開 {{runtime.access_token}} 模板。
//
// Actions:
//   - authenticate: 用 refresh_token 換 access_token，回傳注入後的 headers/query/body
//   - needs_refresh: 檢查 access_token 是否快過期（expires_at < now+300s）
//   - refresh: 強制重新 refresh，更新 KV 中的 cached_access_token/expires_at
//
// Host imports:
//   - u6u.kv_get       — 讀 RECIPES + CREDENTIALS_KV
//   - u6u.kv_put       — 寫回 cached_access_token/expires_at
//   - u6u.crypto_decrypt — AES-GCM 解密 refresh_token
//   - u6u.http_request — POST token endpoint
//
//go:build tinygo

package main

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

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

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

//go:wasmimport u6u kv_put
func hostKvPut(
	keyPtr uintptr, keyLen uint32,
	valPtr uintptr, valLen uint32,
	ttlSeconds uint32,
) uint32

//go:wasmimport u6u crypto_decrypt
func hostCryptoDecrypt(
	encPtr uintptr, encLen uint32,
	ivPtr uintptr, ivLen 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 OAuth2Config struct {
	TokenEndpoint string   `json:"token_endpoint"`
	ClientID      string   `json:"client_id"`
	ClientSecret  string   `json:"client_secret"`
	Scopes        []string `json:"scopes,omitempty"`
}

type AuthRecipe struct {
	Kind            string              `json:"kind"`
	Service         string              `json:"service"`
	Primitive       string              `json:"primitive"`
	OAuth2          *OAuth2Config       `json:"oauth2,omitempty"`
	RequiredSecrets []SecretRequirement `json:"required_secrets"`
	Inject          AuthInjectSpec      `json:"inject"`
}

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

type TokenResponse struct {
	AccessToken  string `json:"access_token"`
	ExpiresIn    int    `json:"expires_in"`
	RefreshToken string `json:"refresh_token,omitempty"`
	Error        string `json:"error"`
	ErrorDesc    string `json:"error_description"`
}

// ── 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
	}

	action := input.Action
	if action == "" {
		action = "authenticate"
	}

	// 讀 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 != "oauth2" {
		writeError("auth recipe " + input.Service + " 的 primitive 不是 oauth2（是 " + recipe.Primitive + "）")
		return
	}
	if recipe.OAuth2 == nil || recipe.OAuth2.TokenEndpoint == "" {
		writeError("auth recipe " + input.Service + " 缺少 oauth2.token_endpoint")
		return
	}

	switch action {
	case "needs_refresh":
		handleNeedsRefresh(input, recipe)
	case "refresh":
		handleRefresh(input, recipe)
	case "authenticate":
		handleAuthenticate(input, recipe)
	default:
		writeError("不支援的 action: " + action)
	}
}

// ── action handlers ───────────────────────────────────────────────────────────

func handleNeedsRefresh(input Input, recipe AuthRecipe) {
	// 讀 cached expires_at（若無，視為需要 refresh）
	expiresKey := input.APIKey + ":oauth2:" + input.Service + ":expires_at"
	expiresStr, status := kvGet(expiresKey)
	if status != 0 {
		// 找不到 = 需要 refresh
		out, _ := json.Marshal(map[string]interface{}{
			"success":      true,
			"needs_refresh": true,
		})
		os.Stdout.Write(out)
		return
	}

	expiresAt, err := strconv.ParseInt(strings.TrimSpace(expiresStr), 10, 64)
	if err != nil {
		out, _ := json.Marshal(map[string]interface{}{
			"success":      true,
			"needs_refresh": true,
		})
		os.Stdout.Write(out)
		return
	}

	// 提前 5 分鐘視為需要 refresh
	needsRefresh := time.Now().Unix()+300 >= expiresAt
	out, _ := json.Marshal(map[string]interface{}{
		"success":      true,
		"needs_refresh": needsRefresh,
	})
	os.Stdout.Write(out)
}

func handleRefresh(input Input, recipe AuthRecipe) {
	accessToken, expiresAt, ok := doRefresh(input, recipe)
	if !ok {
		return
	}

	// 快取新 token
	cacheAccessToken(input.APIKey, input.Service, accessToken, expiresAt)

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

func handleAuthenticate(input Input, recipe AuthRecipe) {
	// 先嘗試讀 cached access_token
	cachedKey := input.APIKey + ":oauth2:" + input.Service + ":access_token"
	expiresKey := input.APIKey + ":oauth2:" + input.Service + ":expires_at"

	cachedToken, cStatus := kvGet(cachedKey)
	expiresStr, eStatus := kvGet(expiresKey)

	var accessToken string
	var expiresAt int64

	useCache := false
	if cStatus == 0 && eStatus == 0 && cachedToken != "" {
		if exp, err := strconv.ParseInt(strings.TrimSpace(expiresStr), 10, 64); err == nil {
			if time.Now().Unix()+300 < exp {
				accessToken = cachedToken
				expiresAt = exp
				useCache = true
			}
		}
	}

	if !useCache {
		var ok bool
		accessToken, expiresAt, ok = doRefresh(input, recipe)
		if !ok {
			return
		}
		cacheAccessToken(input.APIKey, input.Service, accessToken, expiresAt)
	}

	runtime := map[string]string{"access_token": accessToken}
	secrets := map[string]string{} // oauth2 inject 只用 runtime.*，不用 secret.*
	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)
}

// doRefresh 解密 refresh_token，打 token endpoint，回傳 (access_token, expires_at_unix, ok)
func doRefresh(input Input, recipe AuthRecipe) (string, int64, bool) {
	if len(recipe.RequiredSecrets) == 0 {
		writeError("auth recipe " + input.Service + " 缺少 required_secrets（需要 refresh_token）")
		return "", 0, false
	}

	// 慣例：required_secrets[0] 是 refresh_token
	rtReq := recipe.RequiredSecrets[0]
	kvKey := input.APIKey + ":cred:" + rtReq.Key
	encJSON, s := kvGet(kvKey)
	if s == 2 {
		writeError("缺少 credential: " + rtReq.Key + "（" + rtReq.Label + "）。執行 acr creds push 推送")
		return "", 0, false
	}
	if s != 0 {
		writeError("kv_get 失敗(credential " + rtReq.Key + ")")
		return "", 0, false
	}

	var rec EncryptedRecord
	if err := json.Unmarshal([]byte(encJSON), &rec); err != nil {
		writeError("credential " + rtReq.Key + " 格式錯誤: " + err.Error())
		return "", 0, false
	}

	refreshToken, ok := cryptoDecrypt(rec.Encrypted, rec.IV)
	if !ok {
		writeError("credential " + rtReq.Key + " 解密失敗")
		return "", 0, false
	}

	// POST token endpoint（form-urlencoded）
	cfg := recipe.OAuth2
	form := url.Values{}
	form.Set("grant_type", "refresh_token")
	form.Set("refresh_token", refreshToken)
	form.Set("client_id", cfg.ClientID)
	form.Set("client_secret", cfg.ClientSecret)
	formBody := form.Encode()

	headersJSON := `{"Content-Type":"application/x-www-form-urlencoded"}`
	respStr, ok2 := httpRequest(cfg.TokenEndpoint, "POST", headersJSON, formBody)
	if !ok2 {
		writeError("token endpoint HTTP 請求失敗")
		return "", 0, false
	}

	var tokenResp TokenResponse
	if err := json.Unmarshal([]byte(respStr), &tokenResp); err != nil {
		writeError("token endpoint 回應解析失敗: " + err.Error())
		return "", 0, false
	}
	if tokenResp.AccessToken == "" {
		msg := tokenResp.Error
		if tokenResp.ErrorDesc != "" {
			msg += ": " + tokenResp.ErrorDesc
		}
		if msg == "" {
			msg = "access_token 為空"
		}
		writeError("token exchange 失敗: " + msg)
		return "", 0, false
	}

	expiresIn := tokenResp.ExpiresIn
	if expiresIn <= 0 {
		expiresIn = 3600
	}
	expiresAt := time.Now().Unix() + int64(expiresIn)

	// 若 token endpoint 回傳新的 refresh_token（部分服務會 rotate），更新快取
	// 注意：寫回加密 KV 需要 host 支援；此處只快取 access_token（明文短效）
	// refresh_token 的 rotation 需要 kv_put_encrypted host function（未來擴充）

	return tokenResp.AccessToken, expiresAt, true
}

func cacheAccessToken(apiKey, service, accessToken string, expiresAt int64) {
	// 快取明文 access_token（短效，TTL = expires_at - now + 60s buffer）
	ttl := uint32(expiresAt - time.Now().Unix() + 60)
	if ttl > 7200 {
		ttl = 7200
	}

	cachedKey := apiKey + ":oauth2:" + service + ":access_token"
	expiresKey := apiKey + ":oauth2:" + service + ":expires_at"
	expiresStr := strconv.FormatInt(expiresAt, 10)

	kvPut(cachedKey, accessToken, ttl)
	kvPut(expiresKey, expiresStr, ttl)
}

// ── 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)
}

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 kvPut(key, value string, ttlSeconds uint32) {
	keyBytes := []byte(key)
	valBytes := []byte(value)
	if len(keyBytes) == 0 || len(valBytes) == 0 {
		return
	}
	hostKvPut(
		uintptr(unsafe.Pointer(&keyBytes[0])), uint32(len(keyBytes)),
		uintptr(unsafe.Pointer(&valBytes[0])), uint32(len(valBytes)),
		ttlSeconds,
	)
}

func cryptoDecrypt(encB64, ivB64 string) (string, bool) {
	encBytes := []byte(encB64)
	ivBytes := []byte(ivB64)
	if len(encBytes) == 0 || len(ivBytes) == 0 {
		return "", false
	}
	outBuf := make([]byte, 65536)
	var outLen uint32

	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
}

func httpRequest(reqURL, method, headersJSON, body string) (string, bool) {
	urlBytes := []byte(reqURL)
	methodBytes := []byte(method)
	headersBytes := []byte(headersJSON)
	bodyBytes := []byte(body)

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

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

	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
}

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
}