package tools

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"image"
	"image/color"
	"image/png"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"codey/internal/llm"
)

func writePNG(t *testing.T, dir, name string, w, h int) string {
	t.Helper()
	img := image.NewRGBA(image.Rect(0, 0, w, h))
	for x := 0; x < w; x++ {
		for y := 0; y < h; y++ {
			img.Set(x, y, color.RGBA{R: uint8(x), G: uint8(y), B: 128, A: 255})
		}
	}
	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		t.Fatal(err)
	}
	p := filepath.Join(dir, name)
	if err := os.WriteFile(p, buf.Bytes(), 0o644); err != nil {
		t.Fatal(err)
	}
	return p
}

func readBlocks(t *testing.T, tool Tool, path string) (Output, error) {
	t.Helper()
	bt, ok := tool.(BlockTool)
	if !ok {
		t.Fatalf("read_file does not implement BlockTool; an image would go to the model as bytes-as-text")
	}
	raw, _ := json.Marshal(ReadArgs{Path: path})
	return bt.ExecuteBlocks(context.Background(), raw)
}

// A screenshot the agent took is only useful if the model can see it. A PNG
// comes back as an image block with a caption, not as a page of binary.
func TestReadFileReturnsAnImageForAPNG(t *testing.T) {
	dir := t.TempDir()
	writePNG(t, dir, "shot.png", 12, 7)
	track := &fakeTracker{}

	out, err := readBlocks(t, NewRead(dir, track), "shot.png")
	if err != nil {
		t.Fatalf("read: %v", err)
	}
	if len(out.Blocks) != 2 || out.Blocks[0].Type != llm.BlockText || out.Blocks[1].Type != llm.BlockImage {
		t.Fatalf("blocks = %+v, want a caption then an image", out.Blocks)
	}
	img := out.Blocks[1]
	if img.Source == nil || img.Source.MediaType != "image/png" || img.Source.Type != "base64" {
		t.Errorf("image source = %+v, want base64 image/png", img.Source)
	}
	decoded, err := base64.StdEncoding.DecodeString(img.Source.Data)
	if err != nil {
		t.Fatalf("image data is not base64: %v", err)
	}
	if cfg, err := png.DecodeConfig(bytes.NewReader(decoded)); err != nil || cfg.Width != 12 || cfg.Height != 7 {
		t.Errorf("decoded image is not the file that was written: %v %+v", err, cfg)
	}
	for _, want := range []string{"shot.png", "PNG", "12×7"} {
		if !strings.Contains(out.Text, want) {
			t.Errorf("caption %q does not say %q", out.Text, want)
		}
	}
	if out.Text != out.Blocks[0].Text {
		t.Errorf("Text %q should be the caption block %q", out.Text, out.Blocks[0].Text)
	}
	if track.n != 1 {
		t.Errorf("the read was not tracked")
	}
}

func TestReadFileImageMediaTypeFollowsTheExtension(t *testing.T) {
	for ext, want := range map[string]string{".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".gif": "image/gif", ".webp": "image/webp", ".PNG": "image/png"} {
		if got, ok := imageMediaType("x" + ext); !ok || got != want {
			t.Errorf("%s: got %q %v, want %q", ext, got, ok, want)
		}
	}
	for _, ext := range []string{".go", ".svg", ".txt", ""} {
		if _, ok := imageMediaType("x" + ext); ok {
			t.Errorf("%q should not be treated as an image", ext)
		}
	}
}

// Text files keep their line-numbered form: the block path is for pictures.
func TestReadFileTextGoesThroughUnchanged(t *testing.T) {
	dir := t.TempDir()
	writeFile(t, dir, "a.go", "package a\n")
	out, err := readBlocks(t, NewRead(dir, nil), "a.go")
	if err != nil {
		t.Fatal(err)
	}
	if len(out.Blocks) != 0 || !strings.Contains(out.Text, "1\tpackage a") {
		t.Errorf("text read changed shape: %+v", out)
	}
}

func TestReadFileRefusesAnOversizedImage(t *testing.T) {
	dir := t.TempDir()
	p := filepath.Join(dir, "huge.png")
	if err := os.WriteFile(p, make([]byte, maxImageBytes+1), 0o644); err != nil {
		t.Fatal(err)
	}
	_, err := readBlocks(t, NewRead(dir, nil), "huge.png")
	if err == nil || !strings.Contains(err.Error(), "too large") {
		t.Errorf("err = %v, want a size refusal that says what to do instead", err)
	}
}

// The registry hands a BlockTool's blocks through and logs that it did.
func TestRegistryExecutePassesBlocksThrough(t *testing.T) {
	dir := t.TempDir()
	writePNG(t, dir, "s.png", 2, 2)
	r := NewRegistry(nil).Register(NewRead(dir, nil))
	raw, _ := json.Marshal(ReadArgs{Path: "s.png"})
	out, err := r.Execute(context.Background(), "read_file", raw)
	if err != nil {
		t.Fatal(err)
	}
	if len(out.Blocks) != 2 {
		t.Errorf("registry dropped the blocks: %+v", out)
	}
}