Refactor: pngmeta rewrite

2025-03-10 20:08:07 +03:00
parent 10f0efbb2a
commit 0497cdcfe6
5 changed files with 410 additions and 111 deletions
--- a/bot.go
+++ b/bot.go
@@ -94,9 +94,11 @@ func fetchModelName() *models.LLMModels {
 	return &llmModel
 }

+// nolint
 func fetchDSBalance() *models.DSBalance {
 	url := "https://api.deepseek.com/user/balance"
 	method := "GET"
+	// nolint
 	req, err := http.NewRequest(method, url, nil)
 	if err != nil {
 		logger.Warn("failed to create request", "error", err)
@@ -119,6 +121,7 @@ func fetchDSBalance() *models.DSBalance {

 func sendMsgToLLM(body io.Reader) {
 	choseChunkParser()
+	// nolint
 	req, err := http.NewRequest("POST", cfg.CurrentAPI, body)
 	req.Header.Add("Accept", "application/json")
 	req.Header.Add("Content-Type", "application/json")
--- a/llm.go
+++ b/llm.go
@@ -197,8 +197,7 @@ func (ds DeepSeeker) FormMsg(msg, role string, resume bool) (io.Reader, error) {
 	}
 	logger.Debug("checking prompt for /completion", "tool_use", cfg.ToolUse,
 		"msg", msg, "resume", resume, "prompt", prompt)
-	var payload any
-	payload = models.NewDSCompletionReq(prompt, chatBody.Model,
+	payload := models.NewDSCompletionReq(prompt, chatBody.Model,
 		defaultLCPProps["temp"], cfg)
 	data, err := json.Marshal(payload)
 	if err != nil {
--- a/pngmeta/altwriter.go
+++ b/pngmeta/altwriter.go
@@ -0,0 +1,137 @@
+package pngmeta
+
+import (
+	"bytes"
+	"elefant/models"
+	"encoding/base64"
+	"encoding/binary"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"hash/crc32"
+	"io"
+	"os"
+)
+
+const (
+	pngHeader     = "\x89PNG\r\n\x1a\n"
+	textChunkType = "tEXt"
+)
+
+// // PngEmbed holds the key-value pair to embed into the PNG file.
+// type PngEmbed struct {
+// 	Key   string
+// 	Value string
+// }
+
+// WriteToPng embeds the metadata into the specified PNG file and writes the result to outfile.
+func WriteToPng(metadata *models.CharCardSpec, sourcePath, outfile string) error {
+	pngData, err := os.ReadFile(sourcePath)
+	if err != nil {
+		return err
+	}
+
+	jsonData, err := json.Marshal(metadata)
+	if err != nil {
+		return err
+	}
+
+	base64Data := base64.StdEncoding.EncodeToString(jsonData)
+	embedData := PngEmbed{
+		Key:   "elefant", // Replace with appropriate key constant
+		Value: base64Data,
+	}
+
+	var outputBuffer bytes.Buffer
+	if _, err := outputBuffer.Write([]byte(pngHeader)); err != nil {
+		return err
+	}
+
+	chunks, iend, err := processChunks(pngData[8:])
+	if err != nil {
+		return err
+	}
+
+	for _, chunk := range chunks {
+		outputBuffer.Write(chunk)
+	}
+
+	newChunk := createTextChunk(embedData)
+	outputBuffer.Write(newChunk)
+	outputBuffer.Write(iend)
+
+	return os.WriteFile(outfile, outputBuffer.Bytes(), 0666)
+}
+
+// processChunks extracts non-tEXt chunks and locates the IEND chunk
+func processChunks(data []byte) ([][]byte, []byte, error) {
+	var (
+		chunks    [][]byte
+		iendChunk []byte
+		reader    = bytes.NewReader(data)
+	)
+
+	for {
+		var chunkLength uint32
+		if err := binary.Read(reader, binary.BigEndian, &chunkLength); err != nil {
+			if errors.Is(err, io.EOF) {
+				break
+			}
+			return nil, nil, fmt.Errorf("error reading chunk length: %w", err)
+		}
+
+		chunkType := make([]byte, 4)
+		if _, err := reader.Read(chunkType); err != nil {
+			return nil, nil, fmt.Errorf("error reading chunk type: %w", err)
+		}
+
+		chunkData := make([]byte, chunkLength)
+		if _, err := reader.Read(chunkData); err != nil {
+			return nil, nil, fmt.Errorf("error reading chunk data: %w", err)
+		}
+
+		crc := make([]byte, 4)
+		if _, err := reader.Read(crc); err != nil {
+			return nil, nil, fmt.Errorf("error reading CRC: %w", err)
+		}
+
+		fullChunk := bytes.NewBuffer(nil)
+		binary.Write(fullChunk, binary.BigEndian, chunkLength)
+		fullChunk.Write(chunkType)
+		fullChunk.Write(chunkData)
+		fullChunk.Write(crc)
+
+		switch string(chunkType) {
+		case "IEND":
+			iendChunk = fullChunk.Bytes()
+			return chunks, iendChunk, nil
+		case textChunkType:
+			continue // Skip existing tEXt chunks
+		default:
+			chunks = append(chunks, fullChunk.Bytes())
+		}
+	}
+
+	return nil, nil, errors.New("IEND chunk not found")
+}
+
+// createTextChunk generates a valid tEXt chunk with proper CRC
+func createTextChunk(embed PngEmbed) []byte {
+	content := bytes.NewBuffer(nil)
+	content.WriteString(embed.Key)
+	content.WriteByte(0) // Null separator
+	content.WriteString(embed.Value)
+
+	data := content.Bytes()
+	crc := crc32.NewIEEE()
+	crc.Write([]byte(textChunkType))
+	crc.Write(data)
+
+	chunk := bytes.NewBuffer(nil)
+	binary.Write(chunk, binary.BigEndian, uint32(len(data)))
+	chunk.Write([]byte(textChunkType))
+	chunk.Write(data)
+	binary.Write(chunk, binary.BigEndian, crc.Sum32())
+
+	return chunk.Bytes()
+}
--- a/pngmeta/metareader_test.go
+++ b/pngmeta/metareader_test.go
@@ -1,7 +1,19 @@
 package pngmeta

 import (
+	"bytes"
+	"elefant/models"
+	"encoding/base64"
+	"encoding/binary"
+	"encoding/json"
+	"errors"
 	"fmt"
+	"image"
+	"image/color"
+	"image/png"
+	"io"
+	"os"
+	"path/filepath"
 	"testing"
 )

@@ -28,3 +40,155 @@ func TestReadMeta(t *testing.T) {
 		})
 	}
 }
+
+// Test helper: Create a simple PNG image with test shapes
+func createTestImage(t *testing.T) string {
+	img := image.NewRGBA(image.Rect(0, 0, 200, 200))
+	// Fill background with white
+	for y := 0; y < 200; y++ {
+		for x := 0; x < 200; x++ {
+			img.Set(x, y, color.White)
+		}
+	}
+	// Draw a red square
+	for y := 50; y < 150; y++ {
+		for x := 50; x < 150; x++ {
+			img.Set(x, y, color.RGBA{R: 255, A: 255})
+		}
+	}
+	// Draw a blue circle
+	center := image.Point{100, 100}
+	radius := 40
+	for y := center.Y - radius; y <= center.Y+radius; y++ {
+		for x := center.X - radius; x <= center.X+radius; x++ {
+			dx := x - center.X
+			dy := y - center.Y
+			if dx*dx+dy*dy <= radius*radius {
+				img.Set(x, y, color.RGBA{B: 255, A: 255})
+			}
+		}
+	}
+	// Create temp file
+	tmpDir := t.TempDir()
+	fpath := filepath.Join(tmpDir, "test-image.png")
+	f, err := os.Create(fpath)
+	if err != nil {
+		t.Fatalf("Error creating temp file: %v", err)
+	}
+	defer f.Close()
+	if err := png.Encode(f, img); err != nil {
+		t.Fatalf("Error encoding PNG: %v", err)
+	}
+	return fpath
+}
+
+func TestWriteToPng(t *testing.T) {
+	// Create test image
+	srcPath := createTestImage(t)
+	dstPath := filepath.Join(filepath.Dir(srcPath), "output.png")
+	// dstPath := "test.png"
+	// Create test metadata
+	metadata := &models.CharCardSpec{
+		Description: "Test image containing a red square and blue circle on white background",
+	}
+	// Embed metadata
+	if err := WriteToPng(metadata, srcPath, dstPath); err != nil {
+		t.Fatalf("WriteToPng failed: %v", err)
+	}
+	// Verify output file exists
+	if _, err := os.Stat(dstPath); os.IsNotExist(err) {
+		t.Fatalf("Output file not created: %v", err)
+	}
+	// Read and verify metadata
+	t.Run("VerifyMetadata", func(t *testing.T) {
+		data, err := os.ReadFile(dstPath)
+		if err != nil {
+			t.Fatalf("Error reading output file: %v", err)
+		}
+		// Verify PNG header
+		if string(data[:8]) != pngHeader {
+			t.Errorf("Invalid PNG header")
+		}
+		// Extract metadata
+		embedded := extractMetadata(t, data)
+		if embedded.Description != metadata.Description {
+			t.Errorf("Metadata mismatch\nWant: %q\nGot:  %q",
+				metadata.Description, embedded.Description)
+		}
+	})
+	// Optional: Add cleanup if needed
+	// t.Cleanup(func() {
+	// 	os.Remove(dstPath)
+	// })
+}
+
+// Helper to extract embedded metadata from PNG bytes
+func extractMetadata(t *testing.T, data []byte) *models.CharCardSpec {
+	r := bytes.NewReader(data[8:]) // Skip PNG header
+	for {
+		var length uint32
+		if err := binary.Read(r, binary.BigEndian, &length); err != nil {
+			if errors.Is(err, io.EOF) {
+				break
+			}
+			t.Fatalf("Error reading chunk length: %v", err)
+		}
+		chunkType := make([]byte, 4)
+		if _, err := r.Read(chunkType); err != nil {
+			t.Fatalf("Error reading chunk type: %v", err)
+		}
+		// Read chunk data
+		chunkData := make([]byte, length)
+		if _, err := r.Read(chunkData); err != nil {
+			t.Fatalf("Error reading chunk data: %v", err)
+		}
+		// Read and discard CRC
+		if _, err := r.Read(make([]byte, 4)); err != nil {
+			t.Fatalf("Error reading CRC: %v", err)
+		}
+		if string(chunkType) == embType {
+			parts := bytes.SplitN(chunkData, []byte{0}, 2)
+			if len(parts) != 2 {
+				t.Fatalf("Invalid tEXt chunk format")
+			}
+			decoded, err := base64.StdEncoding.DecodeString(string(parts[1]))
+			if err != nil {
+				t.Fatalf("Base64 decode error: %v", err)
+			}
+			var result models.CharCardSpec
+			if err := json.Unmarshal(decoded, &result); err != nil {
+				t.Fatalf("JSON unmarshal error: %v", err)
+			}
+			return &result
+		}
+	}
+	t.Fatal("Metadata not found in PNG")
+	return nil
+}
+
+func readTextChunk(t *testing.T, r io.ReadSeeker) *models.CharCardSpec {
+	var length uint32
+	binary.Read(r, binary.BigEndian, &length)
+	chunkType := make([]byte, 4)
+	r.Read(chunkType)
+	data := make([]byte, length)
+	r.Read(data)
+	// Read CRC (but skip validation for test purposes)
+	crc := make([]byte, 4)
+	r.Read(crc)
+	parts := bytes.SplitN(data, []byte{0}, 2) // Split key-value pair
+	if len(parts) != 2 {
+		t.Fatalf("Invalid tEXt chunk format")
+	}
+	// key := string(parts[0])
+	value := parts[1]
+	decoded, err := base64.StdEncoding.DecodeString(string(value))
+	if err != nil {
+		t.Fatalf("Base64 decode error: %v; value: %s", err, string(value))
+	}
+	var result models.CharCardSpec
+	if err := json.Unmarshal(decoded, &result); err != nil {
+		t.Fatalf("JSON unmarshal error: %v", err)
+	}
+	return &result
+}
--- a/pngmeta/partswriter.go
+++ b/pngmeta/partswriter.go
@@ -1,116 +1,112 @@
 package pngmeta

-import (
-	"bytes"
-	"elefant/models"
-	"encoding/base64"
-	"encoding/binary"
-	"encoding/json"
-	"errors"
-	"fmt"
-	"hash/crc32"
-	"io"
-	"os"
-)
+// import (
+// 	"bytes"
+// 	"encoding/binary"
+// 	"errors"
+// 	"fmt"
+// 	"hash/crc32"
+// 	"io"
+// )

-type Writer struct {
-	w io.Writer
-}
+// type Writer struct {
+// 	w io.Writer
+// }

-func NewPNGWriter(w io.Writer) (*Writer, error) {
-	if _, err := io.WriteString(w, writeHeader); err != nil {
-		return nil, err
-	}
-	return &Writer{w}, nil
-}
+// func NewPNGWriter(w io.Writer) (*Writer, error) {
+// 	if _, err := io.WriteString(w, writeHeader); err != nil {
+// 		return nil, err
+// 	}
+// 	return &Writer{w}, nil
+// }

-func (w *Writer) WriteChunk(length int32, typ string, r io.Reader) error {
-	if err := binary.Write(w.w, binary.BigEndian, length); err != nil {
-		return err
-	}
-	if _, err := w.w.Write([]byte(typ)); err != nil {
-		return err
-	}
-	checksummer := crc32.NewIEEE()
-	checksummer.Write([]byte(typ))
-	if _, err := io.CopyN(io.MultiWriter(w.w, checksummer), r, int64(length)); err != nil {
-		return err
-	}
-	if err := binary.Write(w.w, binary.BigEndian, checksummer.Sum32()); err != nil {
-		return err
-	}
-	return nil
-}
+// func (w *Writer) WriteChunk(length int32, typ string, r io.Reader) error {
+// 	if err := binary.Write(w.w, binary.BigEndian, length); err != nil {
+// 		return err
+// 	}
+// 	if _, err := w.w.Write([]byte(typ)); err != nil {
+// 		return err
+// 	}
+// 	checksummer := crc32.NewIEEE()
+// 	checksummer.Write([]byte(typ))
+// 	if _, err := io.CopyN(io.MultiWriter(w.w, checksummer), r, int64(length)); err != nil {
+// 		return err
+// 	}
+// 	if err := binary.Write(w.w, binary.BigEndian, checksummer.Sum32()); err != nil {
+// 		return err
+// 	}
+// 	return nil
+// }

-func WriteToPng(c *models.CharCardSpec, fpath, outfile string) error {
-	data, err := os.ReadFile(fpath)
-	if err != nil {
-		return err
-	}
-	jsonData, err := json.Marshal(c)
-	if err != nil {
-		return err
-	}
-	// Base64 encode the JSON data
-	base64Data := base64.StdEncoding.EncodeToString(jsonData)
-	pe := PngEmbed{
-		Key:   cKey,
-		Value: base64Data,
-	}
-	w, err := WritetEXtToPngBytes(data, pe)
-	if err != nil {
-		return err
-	}
-	return os.WriteFile(outfile, w.Bytes(), 0666)
-}
+// func WWriteToPngriteToPng(c *models.CharCardSpec, fpath, outfile string) error {
+// 	data, err := os.ReadFile(fpath)
+// 	if err != nil {
+// 		return err
+// 	}
+// 	jsonData, err := json.Marshal(c)
+// 	if err != nil {
+// 		return err
+// 	}
+// 	// Base64 encode the JSON data
+// 	base64Data := base64.StdEncoding.EncodeToString(jsonData)
+// 	pe := PngEmbed{
+// 		Key:   cKey,
+// 		Value: base64Data,
+// 	}
+// 	w, err := WritetEXtToPngBytes(data, pe)
+// 	if err != nil {
+// 		return err
+// 	}
+// 	return os.WriteFile(outfile, w.Bytes(), 0666)
+// }

-func WritetEXtToPngBytes(inputBytes []byte, pe PngEmbed) (outputBytes bytes.Buffer, err error) {
-	if !(string(inputBytes[:8]) == header) {
-		return outputBytes, errors.New("wrong file format")
-	}
-	reader := bytes.NewReader(inputBytes)
-	pngr, err := NewPNGStepReader(reader)
-	if err != nil {
-		return outputBytes, fmt.Errorf("NewReader(): %s", err)
-	}
-	pngw, err := NewPNGWriter(&outputBytes)
-	if err != nil {
-		return outputBytes, fmt.Errorf("NewWriter(): %s", err)
-	}
-	for {
-		chunk, err := pngr.Next()
-		if err != nil {
-			if errors.Is(err, io.EOF) {
-				break
-			}
-			return outputBytes, fmt.Errorf("NextChunk(): %s", err)
-		}
-		if chunk.Type() != embType {
-			// IENDChunkType will only appear on the final iteration of a valid PNG
-			if chunk.Type() == IEND {
-				// This is where we inject tEXtChunkType as the penultimate chunk with the new value
-				newtEXtChunk := []byte(fmt.Sprintf(tEXtChunkDataSpecification, pe.Key, pe.Value))
-				if err := pngw.WriteChunk(int32(len(newtEXtChunk)), embType, bytes.NewBuffer(newtEXtChunk)); err != nil {
-					return outputBytes, fmt.Errorf("WriteChunk(): %s", err)
-				}
-				// Now we end the buffer with IENDChunkType chunk
-				if err := pngw.WriteChunk(chunk.length, chunk.Type(), chunk); err != nil {
-					return outputBytes, fmt.Errorf("WriteChunk(): %s", err)
-				}
-			} else {
-				// writes back original chunk to buffer
-				if err := pngw.WriteChunk(chunk.length, chunk.Type(), chunk); err != nil {
-					return outputBytes, fmt.Errorf("WriteChunk(): %s", err)
-				}
-			}
-		} else {
-			if _, err := io.Copy(io.Discard, chunk); err != nil {
-				return outputBytes, fmt.Errorf("io.Copy(io.Discard, chunk): %s", err)
-			}
-		}
-		if err := chunk.Close(); err != nil {
-			return outputBytes, fmt.Errorf("chunk.Close(): %s", err)
-		}
-	}
-	return outputBytes, nil
-}
+// func WritetEXtToPngBytes(inputBytes []byte, pe PngEmbed) (outputBytes bytes.Buffer, err error) {
+// 	if !(string(inputBytes[:8]) == header) {
+// 		return outputBytes, errors.New("wrong file format")
+// 	}
+// 	reader := bytes.NewReader(inputBytes)
+// 	pngr, err := NewPNGStepReader(reader)
+// 	if err != nil {
+// 		return outputBytes, fmt.Errorf("NewReader(): %s", err)
+// 	}
+// 	pngw, err := NewPNGWriter(&outputBytes)
+// 	if err != nil {
+// 		return outputBytes, fmt.Errorf("NewWriter(): %s", err)
+// 	}
+// 	for {
+// 		chunk, err := pngr.Next()
+// 		if err != nil {
+// 			if errors.Is(err, io.EOF) {
+// 				break
+// 			}
+// 			return outputBytes, fmt.Errorf("NextChunk(): %s", err)
+// 		}
+// 		if chunk.Type() != embType {
+// 			// IENDChunkType will only appear on the final iteration of a valid PNG
+// 			if chunk.Type() == IEND {
+// 				// This is where we inject tEXtChunkType as the penultimate chunk with the new value
+// 				newtEXtChunk := []byte(fmt.Sprintf(tEXtChunkDataSpecification, pe.Key, pe.Value))
+// 				if err := pngw.WriteChunk(int32(len(newtEXtChunk)), embType, bytes.NewBuffer(newtEXtChunk)); err != nil {
+// 					return outputBytes, fmt.Errorf("WriteChunk(): %s", err)
+// 				}
+// 				// Now we end the buffer with IENDChunkType chunk
+// 				if err := pngw.WriteChunk(chunk.length, chunk.Type(), chunk); err != nil {
+// 					return outputBytes, fmt.Errorf("WriteChunk(): %s", err)
+// 				}
+// 			} else {
+// 				// writes back original chunk to buffer
+// 				if err := pngw.WriteChunk(chunk.length, chunk.Type(), chunk); err != nil {
+// 					return outputBytes, fmt.Errorf("WriteChunk(): %s", err)
+// 				}
+// 			}
+// 		} else {
+// 			if _, err := io.Copy(io.Discard, chunk); err != nil {
+// 				return outputBytes, fmt.Errorf("io.Copy(io.Discard, chunk): %s", err)
+// 			}
+// 		}
+// 		if err := chunk.Close(); err != nil {
+// 			return outputBytes, fmt.Errorf("chunk.Close(): %s", err)
+// 		}
+// 	}
+// 	return outputBytes, nil
+// }