package main

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"fmt"
	"image"
	"image/draw"
	"image/jpeg"
	"image/png"
	"io"
	"io/fs"
	"log"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/auyer/steganography"
	"golang.org/x/image/webp"
	"gopkg.in/yaml.v3"
)

type Config struct {
	Server struct {
		Port        string `yaml:"port"`
		ImageDir    string `yaml:"image_dir"`
		BasePath    string `yaml:"base_path"`
		CacheImages bool   `yaml:"cache_images"`
	} `yaml:"server"`
	Security struct {
		AESKey string `yaml:"aes_key"` // Base64 encoded 32-byte key
	} `yaml:"security"`
}

func loadConfig(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)
	}

	// Set defaults if not specified
	if cfg.Server.Port == "" {
		cfg.Server.Port = "8219"
	}
	if cfg.Server.ImageDir == "" {
		cfg.Server.ImageDir = "./images"
	}

	if cfg.Server.BasePath == "" {
		cfg.Server.BasePath = "/"
	}

	return nil
}

var (
	decodedImageCache sync.Map // path -> *cachedImage
	aesKey            []byte
	cfg               Config
)

type cachedImage struct {
	img            *image.NRGBA
	mimeType       string
	maxPayloadSize uint32
	size           uint32
}

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

	if cfg.Security.AESKey != "" {
		var err error
		aesKey, err = base64.StdEncoding.DecodeString(cfg.Security.AESKey)
		if err != nil {
			log.Fatalf("Invalid AES key in config: %v", err)
		}
	} else {
		aesKey = make([]byte, 32)
		if _, err := rand.Read(aesKey); err != nil {
			log.Fatalf("Failed to generate AES key: %v", err)
		}
		log.Printf("Generated new AES key (base64): %s", base64.StdEncoding.EncodeToString(aesKey))
	}

	if cfg.Server.CacheImages {
		if err := loadImages(); err != nil {
			log.Printf("Warning: could not preload images: %v", err)
		}
	}
}

func loadImages() error {
	return filepath.WalkDir(cfg.Server.ImageDir, func(path string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}

		if d.IsDir() {
			return nil
		}

		ext := strings.ToLower(filepath.Ext(path))
		switch ext {
		case ".jpg", ".jpeg", ".png", ".webp":
			data, err := os.ReadFile(path)
			size := uint32(len(data))
			if err != nil {
				return err
			}
			// Preload decoded image
			img, mimeType, err := decodeImageData(data, ext)
			if err != nil {
				return fmt.Errorf("error decoding %s: %v", path, err)
			}

			nrgba := image.NewNRGBA(img.Bounds())
			draw.Draw(nrgba, nrgba.Bounds(), img, img.Bounds().Min, draw.Src)

			maxSize := steganography.MaxEncodeSize(nrgba)
			maxPlaintext, err := maxPlaintextSize(maxSize)
			if err != nil {
				log.Printf("Skipping %s: %v", path, err)
				return nil
			}

			decodedImageCache.Store(path, &cachedImage{
				img:            nrgba,
				mimeType:       mimeType,
				maxPayloadSize: maxPlaintext,
				size:           size,
			})
			log.Printf("Cached decoded image: %s", path)
		}
		return nil
	})
}

func decodeImageData(data []byte, ext string) (image.Image, string, error) {
	switch ext {
	case ".jpg", ".jpeg":
		img, err := jpeg.Decode(bytes.NewReader(data))
		return img, "image/jpeg", err
	case ".png":
		img, err := png.Decode(bytes.NewReader(data))
		return img, "image/png", err
	case ".webp":
		img, err := webp.Decode(bytes.NewReader(data))
		return img, "image/webp", err
	default:
		return nil, "", fmt.Errorf("unsupported format: %s", ext)
	}
}

func encryptData(plaintext []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
	}

	nonce := make([]byte, gcm.NonceSize())
	if _, err = io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, err
	}

	return gcm.Seal(nonce, nonce, plaintext, nil), nil
}

func maxPlaintextSize(totalOutputSize uint32) (uint32, error) {
	block, err := aes.NewCipher(aesKey)
	if err != nil {
		return 0, err
	}

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

	overhead := uint32(gcm.NonceSize() + gcm.Overhead())
	if totalOutputSize < overhead {
		return 0, fmt.Errorf("output size too small")
	}

	return totalOutputSize - overhead, nil
}

var bufferPool = sync.Pool{
	New: func() interface{} {
		return new(bytes.Buffer)
	},
}

func embedData(img *image.NRGBA, data []byte) ([]byte, error) {
	buf := bufferPool.Get().(*bytes.Buffer)
	defer func() {
		buf.Reset()
		if buf.Cap() <= 2<<20 { // 2MB cap
			bufferPool.Put(buf)
		}
	}()

	err := steganography.Encode(buf, img, data)
	if err != nil {
		return nil, fmt.Errorf("failed to encode data into image: %w", err)
	}

	result := make([]byte, buf.Len())
	copy(result, buf.Bytes())
	return result, nil
}

func getClientInfo(r *http.Request) string {
	ip := getClientIP(r)
	var headers strings.Builder
	for name, values := range r.Header {
		for _, value := range values {
			headers.WriteString(fmt.Sprintf("%s: %s\n", name, value))
		}
	}
	return fmt.Sprintf("%s\n%s\n%s", time.Now().UTC().Format(time.RFC3339), ip, headers.String())
}

func getClientIP(r *http.Request) string {
	headers := []string{"X-Forwarded-For"}
	for _, h := range headers {
		if ip := r.Header.Get(h); ip != "" {
			parts := strings.Split(ip, ",")
			clientIP := strings.TrimSpace(parts[0])
			if len(parts) > 1 && isTrustedProxy(clientIP) {
				return strings.TrimSpace(parts[1])
			}
			return clientIP
		}
	}

	ip, _, _ := net.SplitHostPort(r.RemoteAddr)
	return ip
}

func isTrustedProxy(ip string) bool {
	return true // Implement actual IP validation if needed
}

func imageHandler(w http.ResponseWriter, r *http.Request) {
	relativePath := strings.TrimPrefix(r.URL.Path, cfg.Server.BasePath)
	requestedPath := filepath.Clean(strings.TrimPrefix(relativePath, "/"))
	imagePath := filepath.Join(cfg.Server.ImageDir, requestedPath)

	log.Printf("Requested image path: %s", imagePath)

	if !strings.HasPrefix(imagePath, filepath.Clean(cfg.Server.ImageDir)+string(os.PathSeparator)) {
		http.Error(w, "Invalid path", http.StatusBadRequest)
		return
	}

	if _, err := os.Stat(imagePath); errors.Is(err, os.ErrNotExist) {
		http.NotFound(w, r)
		return
	}

	var ci *cachedImage
	if cached, ok := decodedImageCache.Load(imagePath); ok {
		ci = cached.(*cachedImage)
	} else {
		data, err := os.ReadFile(imagePath)
		if err != nil {
			http.Error(w, "File read error", http.StatusInternalServerError)
			return
		}

		ext := strings.ToLower(filepath.Ext(imagePath))
		img, mimeType, err := decodeImageData(data, ext)
		if err != nil {
			http.ServeFile(w, r, imagePath)
			return
		}

		nrgba := image.NewNRGBA(img.Bounds())
		draw.Draw(nrgba, nrgba.Bounds(), img, img.Bounds().Min, draw.Src)

		maxSize := steganography.MaxEncodeSize(nrgba)
		maxPlaintext, err := maxPlaintextSize(maxSize)
		if err != nil {
			http.ServeFile(w, r, imagePath)
			return
		}

		ci = &cachedImage{
			img:            nrgba,
			mimeType:       mimeType,
			maxPayloadSize: maxPlaintext,
			size:           uint32(len(data)),
		}
		decodedImageCache.Store(imagePath, ci)
	}

	metadata := getClientInfo(r)
	if len(metadata) > int(ci.maxPayloadSize) {
		metadata = metadata[:ci.maxPayloadSize]
	}

	encryptedData, err := encryptData([]byte(metadata))
	if err != nil {
		http.ServeFile(w, r, imagePath)
		return
	}

	modifiedImage, err := embedData(ci.img, encryptedData)
	if err != nil {
		http.ServeFile(w, r, imagePath)
		return
	}

	w.Header().Set("Content-Type", "image/png")
	w.Header().Set("Content-Length", fmt.Sprint(len(modifiedImage)))
	w.Write(modifiedImage)
}

func main() {
	if _, err := os.Stat(cfg.Server.ImageDir); os.IsNotExist(err) {
		log.Fatalf("Image directory does not exist: %s", cfg.Server.ImageDir)
	}

	basePath := strings.TrimSuffix(cfg.Server.BasePath, "/")
	http.HandleFunc(basePath+"/", imageHandler)

	log.Printf("Starting server on :%s", cfg.Server.Port)
	log.Fatal(http.ListenAndServe(":"+cfg.Server.Port, nil))
}