codeaware_test.go

package tok_test

import (
	"strings"
	"testing"

	"github.com/GrayCodeAI/tok"
)

const goSample = `package main

import "fmt"

// Greeter greets people by name with a friendly, configurable message that
// can be customized. This comment is deliberately long and chatty so that an
// aggressive prose-aware compression stage has redundant filler to remove.
func Greeter(name string, loud bool) string {
	// build the greeting; honestly this could be much simpler but we add words
	msg := fmt.Sprintf("hello there, %s", name)
	if loud {
		msg = strings.ToUpper(msg)
	}
	return msg
}

type Config struct {
	Name string
	Port int
}
`

func TestRegexSymbolProvider_Go(t *testing.T) {
	spans := tok.RegexSymbolProvider{}.Symbols(goSample, "go")
	if len(spans) == 0 {
		t.Fatal("expected spans for Go source, got none")
	}

	want := map[string]bool{"Greeter": false, "Config": false}
	for _, s := range spans {
		if s.StartLine <= 0 || s.EndLine < s.StartLine {
			t.Errorf("invalid span range %d-%d", s.StartLine, s.EndLine)
		}
		if len(s.Lines) == 0 {
			t.Errorf("span %q has no lines", s.Symbol)
		}
		if _, ok := want[s.Symbol]; ok {
			want[s.Symbol] = true
		}
	}
	for sym, found := range want {
		if !found {
			t.Errorf("expected to find symbol %q in spans", sym)
		}
	}
}

func TestRegexSymbolProvider_UnknownLanguage(t *testing.T) {
	if got := (tok.RegexSymbolProvider{}).Symbols("some plain prose here", "cobol-ish"); got != nil {
		t.Errorf("expected nil spans for unknown language, got %v", got)
	}
}

func TestCompress_CodeAwarePreservesSignatures(t *testing.T) {
	signatures := []string{
		"func Greeter(name string, loud bool) string {",
		"type Config struct {",
	}

	out, _ := tok.Compress(goSample, tok.WithCodeAware("go"), tok.Aggressive)

	norm := strings.Join(strings.Fields(out), " ")
	for _, sig := range signatures {
		needle := strings.Join(strings.Fields(sig), " ")
		if !strings.Contains(norm, needle) {
			t.Errorf("code-aware output dropped signature: %q\n---\n%s", sig, out)
		}
	}
}

func TestCompress_ProseStillCompresses(t *testing.T) {
	prose := strings.Repeat(
		"Sure, I would be happy to help you with that particular request, and "+
			"I will absolutely do my very best to assist you in every possible way. ",
		20,
	)

	plain, _ := tok.Compress(prose, tok.Aggressive)
	coded, _ := tok.Compress(prose, tok.WithCodeAware("go"), tok.Aggressive)

	if len(plain) >= len(prose) {
		t.Errorf("expected prose to compress: in=%d out=%d", len(prose), len(plain))
	}
	// Prose has no Go signatures, so the guard must not change the result.
	if coded != plain {
		t.Errorf("code-aware guard altered prose output\nplain: %q\ncoded: %q", plain, coded)
	}
}

// stubProvider is a custom SymbolProvider used to verify the injection seam.
type stubProvider struct{ called bool }

func (s *stubProvider) Symbols(code, lang string) []tok.Span {
	s.called = true
	return []tok.Span{{
		StartLine: 1,
		EndLine:   1,
		Symbol:    "SENTINEL",
		Lines:     []string{"SENTINEL_SIGNATURE_LINE_THAT_DOES_NOT_EXIST"},
	}}
}

func TestWithSymbolProvider_CustomProviderUsedAndGuarded(t *testing.T) {
	sp := &stubProvider{}
	out, _ := tok.Compress(goSample, tok.WithCodeAware("go"), tok.WithSymbolProvider(sp), tok.Aggressive)

	if !sp.called {
		t.Fatal("custom SymbolProvider was not invoked")
	}
	if !strings.Contains(out, "SENTINEL_SIGNATURE_LINE_THAT_DOES_NOT_EXIST") {
		t.Errorf("guard did not re-inject missing protected line from custom provider:\n%s", out)
	}
}

func TestWithSymbolProvider_NilIsNoop(t *testing.T) {
	// Passing nil must not panic and must fall back to the default provider.
	out, _ := tok.Compress(goSample, tok.WithCodeAware("go"), tok.WithSymbolProvider(nil), tok.Aggressive)
	if !strings.Contains(strings.Join(strings.Fields(out), " "), "func Greeter(name string, loud bool) string {") {
		t.Errorf("default provider not used after nil WithSymbolProvider:\n%s", out)
	}
}