package app

import (
	"bytes"
	"encoding/json"
	"fmt"
	"verification-platform-backend/pkg/crypto"

	"github.com/gin-gonic/gin"
)

type bodyLogWriter struct {
	gin.ResponseWriter
	body *bytes.Buffer
}

func (w *bodyLogWriter) Write(b []byte) (int, error) {
	return w.body.Write(b)
}

func (w *bodyLogWriter) WriteHeader(statusCode int) {
	w.ResponseWriter.WriteHeader(statusCode)
}

type EncryptedResponse struct {
	Data string `json:"data"`
}

func ResponseEncryption() gin.HandlerFunc {
	return func(c *gin.Context) {
		blw := &bodyLogWriter{body: bytes.NewBufferString(""), ResponseWriter: c.Writer}
		c.Writer = blw

		c.Next()

		shouldEncrypt := c.GetBool("should_encrypt_response")
		fmt.Printf("[ResponseEncryption] ShouldEncrypt: %v\n", shouldEncrypt)

		blw.ResponseWriter.Header().Set("Content-Type", "application/json; charset=utf-8")

		if shouldEncrypt {
			var response map[string]interface{}
			if err := json.Unmarshal(blw.body.Bytes(), &response); err == nil {
				responseJSON, _ := json.Marshal(response)
				fmt.Printf("[ResponseEncryption] Response to encrypt: %s\n", string(responseJSON))

				var cryptoManager *crypto.CryptoManager
				if cm, ok := c.Get("crypto_manager"); ok {
					cryptoManager = cm.(*crypto.CryptoManager)
				}

				if cryptoManager != nil {
					encrypted, err := cryptoManager.Encrypt(string(responseJSON))
					if err == nil {
						encryptedResponse := EncryptedResponse{Data: encrypted}
						encryptedJSON, _ := json.Marshal(encryptedResponse)
						blw.ResponseWriter.Write(encryptedJSON)
						fmt.Printf("[ResponseEncryption] Encrypted response sent\n")
						return
					} else {
						fmt.Printf("[ResponseEncryption] Encryption error: %v\n", err)
					}
				} else {
					fmt.Printf("[ResponseEncryption] CryptoManager is nil\n")
				}
			} else {
				fmt.Printf("[ResponseEncryption] JSON unmarshal error: %v\n", err)
			}
		}

		blw.ResponseWriter.Write(blw.body.Bytes())
	}
}