package agent

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"strings"
)

type StoredAgentLLMColumns struct {
	Mode             string
	BaseURL          string
	APIKeyCiphertext string
	HeadersJSON      string
	WireAPI          string
}

func HasAgentCredentialStorage(key []byte) bool {
	return len(key) == 32
}

func EncryptAgentCredential(key []byte, plaintext string) (string, error) {
	if !HasAgentCredentialStorage(key) {
		return "", fmt.Errorf("agent credential storage is not configured on server")
	}

	block, err := aes.NewCipher(key)
	if err != nil {
		return "", fmt.Errorf("create agent credential cipher: %w", err)
	}
	aead, err := cipher.NewGCM(block)
	if err != nil {
		return "", fmt.Errorf("create agent credential gcm: %w", err)
	}

	nonce := make([]byte, aead.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		return "", fmt.Errorf("generate agent credential nonce: %w", err)
	}

	sealed := aead.Seal(nonce, nonce, []byte(plaintext), nil)
	return base64.RawStdEncoding.EncodeToString(sealed), nil
}

func DecryptAgentCredential(key []byte, ciphertext string) (string, error) {
	return DecryptAgentCredentialWithFallback([][]byte{key}, ciphertext)
}

func DecryptAgentCredentialWithFallback(keys [][]byte, ciphertext string) (string, error) {
	var firstErr error
	for _, key := range keys {
		if !HasAgentCredentialStorage(key) {
			if firstErr == nil {
				firstErr = fmt.Errorf("agent credential storage is not configured on server")
			}
			continue
		}
		plaintext, err := decryptAgentCredentialWithKey(key, ciphertext)
		if err == nil {
			return plaintext, nil
		}
		if firstErr == nil {
			firstErr = err
		}
	}
	if firstErr != nil {
		return "", firstErr
	}
	return "", fmt.Errorf("agent credential storage is not configured on server")
}

func decryptAgentCredentialWithKey(key []byte, ciphertext string) (string, error) {
	if !HasAgentCredentialStorage(key) {
		return "", fmt.Errorf("agent credential storage is not configured on server")
	}

	block, err := aes.NewCipher(key)
	if err != nil {
		return "", fmt.Errorf("create agent credential cipher: %w", err)
	}
	aead, err := cipher.NewGCM(block)
	if err != nil {
		return "", fmt.Errorf("create agent credential gcm: %w", err)
	}

	payload, err := base64.RawStdEncoding.DecodeString(strings.TrimSpace(ciphertext))
	if err != nil {
		return "", fmt.Errorf("decode agent credential ciphertext: %w", err)
	}
	if len(payload) < aead.NonceSize() {
		return "", fmt.Errorf("agent credential ciphertext is too short")
	}

	nonce := payload[:aead.NonceSize()]
	encrypted := payload[aead.NonceSize():]
	plaintext, err := aead.Open(nil, nonce, encrypted, nil)
	if err != nil {
		return "", fmt.Errorf("decrypt agent credential ciphertext: %w", err)
	}
	return string(plaintext), nil
}

func EncodeStoredAgentLLMColumns(key []byte, llm *AgentLLM) (StoredAgentLLMColumns, error) {
	normalized := llm.Normalize()
	if normalized == nil {
		normalized = (&AgentLLM{Mode: AgentLLMModeLocal}).Normalize()
	}

	headersJSON, err := EncodeAgentHeaders(normalized.Headers)
	if err != nil {
		return StoredAgentLLMColumns{}, err
	}

	apiKeyCiphertext := ""
	if strings.TrimSpace(normalized.APIKey) != "" {
		apiKeyCiphertext, err = EncryptAgentCredential(key, normalized.APIKey)
		if err != nil {
			return StoredAgentLLMColumns{}, err
		}
	}

	return StoredAgentLLMColumns{
		Mode:             normalized.Mode,
		BaseURL:          normalized.BaseURL,
		APIKeyCiphertext: apiKeyCiphertext,
		HeadersJSON:      headersJSON,
		WireAPI:          normalized.WireAPI,
	}, nil
}

func DecodeStoredAgentLLMColumns(key []byte, mode string, baseURL string, apiKeyCiphertext string, headersJSON []byte, wireAPI string) (*AgentLLM, error) {
	return DecodeStoredAgentLLMColumnsWithFallback(key, nil, mode, baseURL, apiKeyCiphertext, headersJSON, wireAPI)
}

func DecodeStoredAgentLLMColumnsWithFallback(key []byte, legacyKeys [][]byte, mode string, baseURL string, apiKeyCiphertext string, headersJSON []byte, wireAPI string) (*AgentLLM, error) {
	headers, err := DecodeAgentHeaders(headersJSON)
	if err != nil {
		return nil, err
	}

	apiKey := ""
	if strings.TrimSpace(apiKeyCiphertext) != "" {
		apiKey, err = DecryptAgentCredentialWithFallback(credentialStorageKeyCandidates(key, legacyKeys), apiKeyCiphertext)
		if err != nil {
			return nil, err
		}
	}

	return (&AgentLLM{
		Mode:    mode,
		BaseURL: baseURL,
		APIKey:  apiKey,
		Headers: headers,
		WireAPI: wireAPI,
	}).Normalize(), nil
}

func cloneCredentialStorageKeys(keys [][]byte) [][]byte {
	if len(keys) == 0 {
		return nil
	}
	cloned := make([][]byte, 0, len(keys))
	for _, key := range keys {
		if len(key) == 0 {
			continue
		}
		cloned = append(cloned, append([]byte(nil), key...))
	}
	return cloned
}

func credentialStorageKeyCandidates(primary []byte, legacy [][]byte) [][]byte {
	candidates := make([][]byte, 0, 1+len(legacy))
	if len(primary) > 0 {
		candidates = append(candidates, primary)
	}
	for _, key := range legacy {
		if len(key) == 0 || credentialStorageKeySeen(candidates, key) {
			continue
		}
		candidates = append(candidates, key)
	}
	return candidates
}

func credentialStorageKeySeen(keys [][]byte, candidate []byte) bool {
	for _, key := range keys {
		if string(key) == string(candidate) {
			return true
		}
	}
	return false
}