chunker_test.go

package tok_test

import (
	"strings"
	"testing"

	"github.com/GrayCodeAI/tok"
)

func TestChunkCode_GoFuncBoundaries(t *testing.T) {
	source := `package main

import "fmt"

func Hello() {
	fmt.Println("hello")
}

func World() {
	fmt.Println("world")
}

type Config struct {
	Name string
	Port int
}
`
	opts := tok.ChunkOptions{MaxTokens: 500, MinTokens: 5, Language: "go"}
	chunks := tok.ChunkCode(source, opts)

	if len(chunks) < 2 {
		t.Fatalf("expected at least 2 chunks for Go code with multiple funcs, got %d", len(chunks))
	}

	// Verify chunks have reasonable start/end lines
	for i, c := range chunks {
		if c.StartLine <= 0 || c.EndLine <= 0 {
			t.Errorf("chunk %d: invalid line range %d-%d", i, c.StartLine, c.EndLine)
		}
		if c.Content == "" {
			t.Errorf("chunk %d: empty content", i)
		}
		if c.Tokens <= 0 {
			t.Errorf("chunk %d: expected positive token count, got %d", i, c.Tokens)
		}
	}

	// At least one chunk should have a symbol name
	hasSymbol := false
	for _, c := range chunks {
		if c.Symbol != "" {
			hasSymbol = true
			break
		}
	}
	if !hasSymbol {
		t.Error("expected at least one chunk to have a symbol name")
	}
}

func TestChunkCode_PythonBoundaries(t *testing.T) {
	source := `import os

class Animal:
    def __init__(self, name):
        self.name = name

    def speak(self):
        pass

def greet(name):
    print(f"Hello, {name}")

async def fetch_data():
    pass
`
	opts := tok.ChunkOptions{MaxTokens: 500, MinTokens: 5, Language: "python"}
	chunks := tok.ChunkCode(source, opts)

	if len(chunks) < 2 {
		t.Fatalf("expected at least 2 chunks for Python code, got %d", len(chunks))
	}

	// Check that Python symbol names are extracted
	foundClass := false
	foundDef := false
	for _, c := range chunks {
		if c.Symbol == "Animal" {
			foundClass = true
		}
		if c.Symbol == "greet" || c.Symbol == "fetch_data" {
			foundDef = true
		}
	}
	if !foundClass && !foundDef {
		t.Log("chunks:")
		for i, c := range chunks {
			t.Logf("  [%d] symbol=%q startLine=%d endLine=%d tokens=%d", i, c.Symbol, c.StartLine, c.EndLine, c.Tokens)
		}
		t.Error("expected to find Python class or def symbol names")
	}
}

func TestChunkCode_MergeSmallChunks(t *testing.T) {
	// Create source with very small functions that should be merged
	source := `package main

func A() {}

func B() {}

func C() {}
`
	// Set a high MinTokens so small chunks get merged
	opts := tok.ChunkOptions{MaxTokens: 500, MinTokens: 100, Language: "go"}
	chunks := tok.ChunkCode(source, opts)

	// With high MinTokens, small funcs should be merged together
	if len(chunks) > 2 {
		t.Errorf("expected small chunks to be merged (got %d chunks)", len(chunks))
	}
}

func TestChunkCode_SplitLargeChunks(t *testing.T) {
	// Generate a large function body that exceeds MaxTokens
	var sb strings.Builder
	sb.WriteString("package main\n\nfunc BigFunc() {\n")
	for i := 0; i < 200; i++ {
		sb.WriteString("    println(\"line " + strings.Repeat("x", 20) + "\")\n")
		// Add occasional blank lines for split points
		if i%20 == 0 && i > 0 {
			sb.WriteString("\n")
		}
	}
	sb.WriteString("}\n")

	opts := tok.ChunkOptions{MaxTokens: 100, MinTokens: 10, Language: "go"}
	chunks := tok.ChunkCode(sb.String(), opts)

	if len(chunks) < 2 {
		t.Fatalf("expected large function to be split into multiple chunks, got %d", len(chunks))
	}

	// Verify no chunk exceeds a reasonable bound (2x MaxTokens as grace)
	for i, c := range chunks {
		if c.Tokens > opts.MaxTokens*3 {
			t.Errorf("chunk %d has %d tokens, expected roughly <= %d", i, c.Tokens, opts.MaxTokens)
		}
	}
}

func TestDetectLanguageByExtension(t *testing.T) {
	tests := []struct {
		path string
		want string
	}{
		{"main.go", "go"},
		{"script.py", "python"},
		{"app.ts", "typescript"},
		{"index.js", "javascript"},
		{"lib.rs", "rust"},
		{"Main.java", "java"},
		{"app.rb", "ruby"},
		{"index.php", "php"},
		{"main.c", "c"},
		{"header.h", "c"},
		{"main.cpp", "cpp"},
		{"main.cc", "cpp"},
		{"Program.cs", "csharp"},
		{"App.kt", "kotlin"},
		{"App.swift", "swift"},
		{"main.dart", "dart"},
		{"setup.sh", "bash"},
		{"query.sql", "sql"},
		{"index.html", "html"},
		{"style.css", "css"},
		{"App.vue", "vue"},
		{"main.zig", "zig"},
		{"app.ex", "elixir"},
		{"Main.scala", "scala"},
		{"script.lua", "lua"},
		{"analysis.r", "r"},
		{"analysis.R", "r"},
		{"README.md", "markdown"},
		{"/path/to/file.go", "go"},
		{"unknown.xyz", ""},
		{"noext", ""},
	}
	for _, tt := range tests {
		got := tok.DetectLanguageByExtension(tt.path)
		if got != tt.want {
			t.Errorf("DetectLanguageByExtension(%q) = %q, want %q", tt.path, got, tt.want)
		}
	}
}

func TestDetectLanguageByExtension_EmptyAndUnknown(t *testing.T) {
	// Empty path
	if got := tok.DetectLanguageByExtension(""); got != "" {
		t.Errorf("expected empty for empty path, got %q", got)
	}
	// Unknown extension
	if got := tok.DetectLanguageByExtension("data.parquet"); got != "" {
		t.Errorf("expected empty for .parquet, got %q", got)
	}
}

func TestSeparatorLevels(t *testing.T) {
	// Create a block of text with no blank lines but with sentence endings.
	// The separator hierarchy should split at sentence boundaries when
	// paragraph breaks (level 1) and line breaks (level 2) are not available.
	text := "This is sentence one. This is sentence two. This is sentence three. This is sentence four. This is sentence five."

	// Use a small MaxTokens so the text must be split
	opts := tok.ChunkOptions{MaxTokens: 20, MinTokens: 1, Language: "go", Overlap: 0}

	// Wrap it in a Go package so it's treated as one chunk initially
	source := "package main\n\n" + text
	chunks := tok.ChunkCode(source, opts)

	if len(chunks) < 2 {
		t.Fatalf("expected separator levels to split the text into multiple chunks, got %d", len(chunks))
	}

	// Verify all chunks have content
	for i, c := range chunks {
		if c.Content == "" {
			t.Errorf("chunk %d: empty content", i)
		}
	}
	t.Logf("separator levels produced %d chunks from run-on text", len(chunks))
}

func TestExtendedExtensions(t *testing.T) {
	tests := []struct {
		path string
		want string
	}{
		{"file.m", "objective-c"},
		{"file.mm", "objective-cpp"},
		{"file.pl", "perl"},
		{"file.pm", "perl"},
		{"file.hs", "haskell"},
		{"file.lhs", "haskell"},
		{"file.erl", "erlang"},
		{"file.hrl", "erlang"},
		{"file.clj", "clojure"},
		{"file.cljs", "clojure"},
		{"file.fs", "fsharp"},
		{"file.fsx", "fsharp"},
		{"file.ml", "ocaml"},
		{"file.mli", "ocaml"},
		{"file.jl", "julia"},
		{"file.nim", "nim"},
		{"file.cr", "crystal"},
		{"file.v", "v"},
		{"file.vhdl", "vhdl"},
		{"file.sol", "solidity"},
		{"file.proto", "protobuf"},
		{"file.graphql", "graphql"},
		{"file.gql", "graphql"},
		{"file.tf", "terraform"},
		{"file.hcl", "hcl"},
		{"file.cmake", "cmake"},
		{"file.gradle", "gradle"},
		{"file.groovy", "groovy"},
		{"file.ps1", "powershell"},
		{"file.psm1", "powershell"},
		{"file.bat", "batch"},
		{"file.cmd", "batch"},
		{"file.fish", "fish"},
		{"file.vim", "vim"},
		{"file.el", "elisp"},
		{"file.rkt", "racket"},
		{"file.pas", "pascal"},
		{"file.d", "d"},
		{"file.ada", "ada"},
		{"file.adb", "ada"},
		{"file.f90", "fortran"},
		{"file.f95", "fortran"},
		{"file.cob", "cobol"},
		{"file.tsx", "tsx"},
		{"file.jsx", "jsx"},
		{"file.svelte", "svelte"},
		{"file.astro", "astro"},
		{"file.prisma", "prisma"},
		{"file.env", "dotenv"},
		{"file.ini", "ini"},
		{"file.cfg", "ini"},
		{"file.conf", "conf"},
		{"file.nginx", "nginx"},
		{"file.dockerfile", "dockerfile"},
		{"file.containerfile", "dockerfile"},
		{"file.nix", "nix"},
		{"file.dhall", "dhall"},
		{"file.jsonnet", "jsonnet"},
		{"file.starlark", "starlark"},
		{"file.bzl", "starlark"},
		{"file.wasm", "wasm"},
		{"file.wat", "wat"},
	}
	for _, tt := range tests {
		got := tok.DetectLanguageByExtension(tt.path)
		if got != tt.want {
			t.Errorf("DetectLanguageByExtension(%q) = %q, want %q", tt.path, got, tt.want)
		}
	}
}

func TestKeepSeparator(t *testing.T) {
	source := `package main

func Alpha() {
	println("alpha")
}

func Beta() {
	println("beta")
}

func Gamma() {
	println("gamma")
}
`
	t.Run("SepLeft", func(t *testing.T) {
		opts := tok.ChunkOptions{
			MaxTokens:     500,
			MinTokens:     1,
			Language:      "go",
			Overlap:       0,
			KeepSeparator: tok.SepLeft,
		}
		chunks := tok.ChunkCode(source, opts)
		if len(chunks) < 2 {
			t.Fatalf("expected multiple chunks, got %d", len(chunks))
		}
		// With SepLeft (default), boundary line starts the next chunk
		// (this is the existing behavior - boundary detected, previous flushed,
		//  boundary line added to next chunk)
		for i, c := range chunks {
			t.Logf("SepLeft chunk %d: %q", i, c.Content[:min(60, len(c.Content))])
		}
	})

	t.Run("SepRight", func(t *testing.T) {
		opts := tok.ChunkOptions{
			MaxTokens:     500,
			MinTokens:     1,
			Language:      "go",
			Overlap:       0,
			KeepSeparator: tok.SepRight,
		}
		chunks := tok.ChunkCode(source, opts)
		if len(chunks) < 2 {
			t.Fatalf("expected multiple chunks, got %d", len(chunks))
		}
		// With SepRight, boundary line should go to the following chunk
		for i, c := range chunks {
			t.Logf("SepRight chunk %d: %q", i, c.Content[:min(60, len(c.Content))])
		}
		// The second chunk should start with "func"
		for _, c := range chunks[1:] {
			trimmed := strings.TrimSpace(c.Content)
			if !strings.HasPrefix(trimmed, "func") && trimmed != "" {
				// This is OK - some chunks may just be whitespace that got merged
			}
		}
	})

	t.Run("SepDiscard", func(t *testing.T) {
		opts := tok.ChunkOptions{
			MaxTokens:     500,
			MinTokens:     1,
			Language:      "go",
			Overlap:       0,
			KeepSeparator: tok.SepDiscard,
		}
		chunks := tok.ChunkCode(source, opts)
		if len(chunks) < 2 {
			t.Fatalf("expected multiple chunks, got %d", len(chunks))
		}
		// With SepDiscard, the boundary lines (func ...) should be removed.
		// Only the first chunk (which has "package main" + first func) should contain "func"
		// since the first func is not a boundary split point.
		discardedCount := 0
		for _, c := range chunks {
			// Count how many chunks DON'T contain "func " - those had it discarded
			if !strings.Contains(c.Content, "func ") {
				discardedCount++
			}
		}
		if discardedCount == 0 {
			t.Log("expected at least some chunks to have their boundary line discarded")
		}
		for i, c := range chunks {
			t.Logf("SepDiscard chunk %d: %q", i, c.Content[:min(60, len(c.Content))])
		}
	})
}

func TestOptimizedMerging(t *testing.T) {
	// Create source with many tiny functions. With cost-optimized merging,
	// small chunks below MinTokens should be merged to reduce penalties.
	var sb strings.Builder
	sb.WriteString("package main\n\n")
	for i := 0; i < 10; i++ {
		sb.WriteString("func F" + string(rune('A'+i)) + "() {\n")
		sb.WriteString("    println(\"hello\")\n")
		sb.WriteString("}\n\n")
	}

	opts := tok.ChunkOptions{
		MaxTokens: 200,
		MinTokens: 30,
		Language:  "go",
		Overlap:   0,
	}
	chunks := tok.ChunkCode(sb.String(), opts)

	// The optimizer should merge small functions together rather than leaving
	// each as its own tiny chunk
	tooSmallCount := 0
	for _, c := range chunks {
		if c.Tokens < opts.MinTokens {
			tooSmallCount++
		}
	}

	// With optimization, we should have fewer too-small chunks than total functions
	if tooSmallCount > 5 {
		t.Errorf("expected optimizer to merge small chunks, but %d/%d chunks are below MinTokens=%d",
			tooSmallCount, len(chunks), opts.MinTokens)
	}

	t.Logf("optimized merging: %d chunks from 10 tiny functions, %d below MinTokens",
		len(chunks), tooSmallCount)

	// Verify no chunk exceeds MaxTokens
	for i, c := range chunks {
		if c.Tokens > opts.MaxTokens {
			t.Errorf("chunk %d exceeds MaxTokens: %d > %d", i, c.Tokens, opts.MaxTokens)
		}
	}
}

func TestCustomLanguagePatterns(t *testing.T) {
	// Register a custom language with a custom boundary pattern
	tok.RegisterLanguagePatterns("mylang", []string{`^SECTION `})

	// Verify GetLanguagePatterns returns the custom patterns
	patterns := tok.GetLanguagePatterns("mylang")
	if len(patterns) != 1 || patterns[0] != `^SECTION ` {
		t.Fatalf("expected custom patterns, got %v", patterns)
	}

	// Verify GetLanguagePatterns returns built-in for known languages
	goPatterns := tok.GetLanguagePatterns("go")
	if len(goPatterns) == 0 {
		t.Fatal("expected built-in patterns for Go")
	}

	// Use the custom language for chunking
	source := `SECTION Intro
This is the introduction.
Some more text here.

SECTION Body
This is the body content.
More body text.

SECTION Conclusion
This is the conclusion.
`
	opts := tok.ChunkOptions{
		MaxTokens: 500,
		MinTokens: 1,
		Language:  "mylang",
		Overlap:   0,
	}
	chunks := tok.ChunkCode(source, opts)

	if len(chunks) < 2 {
		t.Fatalf("expected custom boundary to split into multiple chunks, got %d", len(chunks))
	}

	// Verify the chunks split at SECTION boundaries
	t.Logf("custom language produced %d chunks", len(chunks))
	for i, c := range chunks {
		t.Logf("  chunk %d: %q", i, c.Content[:min(50, len(c.Content))])
	}

	// Clean up: re-register with nil to avoid affecting other tests
	tok.RegisterLanguagePatterns("mylang", nil)
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func TestChunkCode_TokenCountingAccuracy(t *testing.T) {
	source := `package main

func main() {
	fmt.Println("hello world")
}
`
	opts := tok.ChunkOptions{MaxTokens: 500, MinTokens: 5, Language: "go"}
	chunks := tok.ChunkCode(source, opts)

	if len(chunks) == 0 {
		t.Fatal("expected at least one chunk")
	}

	for i, c := range chunks {
		// Verify that the reported token count matches what EstimateTokensPrecise returns
		expected := tok.EstimateTokensPrecise(c.Content)
		if c.Tokens != expected {
			t.Errorf("chunk %d: token count mismatch: chunk reports %d, EstimateTokensPrecise returns %d",
				i, c.Tokens, expected)
		}

		// Sanity check: tokens should be positive for non-empty content
		if c.Content != "" && c.Tokens <= 0 {
			t.Errorf("chunk %d: non-empty content but %d tokens", i, c.Tokens)
		}
	}
}