package main

import (
	"bufio"
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"
	"fmt"
	"image"
	"log"
	"os"

	"github.com/auyer/steganography"
	"github.com/chai2010/webp"
	"gopkg.in/yaml.v3"
)

type DecodeConfig struct {
	Security struct {
		AESKey string `yaml:"aes_key"` // Base64 encoded 32-byte key
	} `yaml:"security"`
}

var (
	aesKey []byte
	cfg    DecodeConfig
)

func loadDecodeConfig(filename string) error {
	data, err := os.ReadFile(filename)
	if err != nil {
		return fmt.Errorf("config file read error: %w", err)
	}

	if err := yaml.Unmarshal(data, &cfg); err != nil {
		return fmt.Errorf("config parse error: %w", err)
	}
	return nil
}

func decryptData(ciphertext []byte) ([]byte, error) {
	block, err := aes.NewCipher(aesKey)
	if err != nil {
		return nil, err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}

	nonceSize := gcm.NonceSize()
	if len(ciphertext) < nonceSize {
		return nil, fmt.Errorf("ciphertext too short")
	}

	nonce, ciphertext := ciphertext[:nonceSize], ciphertext[nonceSize:]
	return gcm.Open(nil, nonce, ciphertext, nil)
}

func decodeImage(file *os.File) (image.Image, error) {
	reader := bufio.NewReader(file)

	// Try to decode with image.Decode first
	img, _, err := image.Decode(reader)
	if err == nil {
		return img, nil
	}

	// Try fallback: WebP
	// Reset file offset to 0
	_, _ = file.Seek(0, 0)
	img, err = webp.Decode(file)
	if err != nil {
		return nil, fmt.Errorf("unsupported image format or decode failure: %w", err)
	}
	return img, nil
}

func main() {
	if len(os.Args) < 2 {
		fmt.Println("Usage: go run decode.go <encoded_image>")
		return
	}

	if err := loadDecodeConfig("config.yaml"); err != nil {
		log.Fatalf("Failed to load config: %v", err)
	}

	filePath := os.Args[1]

	inFile, err := os.Open(filePath)
	if err != nil {
		fmt.Printf("Failed to open file: %v\n", err)
		return
	}
	defer inFile.Close()

	// Initialize AES key
	if cfg.Security.AESKey != "" {
		aesKey, err = base64.StdEncoding.DecodeString(cfg.Security.AESKey)
		if err != nil {
			log.Fatalf("Invalid AES key in config: %v", err)
		}
	} else {
		log.Fatal("AES key not provided in config.")
	}

	img, err := decodeImage(inFile)
	if err != nil {
		fmt.Printf("Failed to decode image: %v\n", err)
		return
	}

	size := steganography.GetMessageSizeFromImage(img)
	message := steganography.Decode(size, img)
	if message == nil {
		fmt.Println("No message found in the image.")
		return
	}

	decryptedMessage, err := decryptData(message)
	if err != nil {
		fmt.Printf("Failed to decrypt message: %v\n", err)
		return
	}

	fmt.Printf("Decrypted message: %s\n", decryptedMessage)
}