package tests

import (
	"strconv"

	"github.com/lukaszraczylo/llm-testbench/internal/eval"
	"github.com/lukaszraczylo/llm-testbench/internal/testkit"
)

// codeExactAnswer returns an Evaluator awarding full credit when the
// response equals want under eval.ExactToken's normalization: a fenced
// code block extracted if present, whitespace trimmed, one layer of
// surrounding quote/backtick/asterisk decoration stripped, one trailing
// sentence-ending period stripped, compared case-insensitively. This
// accepts every materially-correct form of a forced single-token answer
// (bare, quoted, bolded, fenced, differently-cased, or with trailing
// punctuation) without loosening the match to accept a substring of a
// longer, wrong answer. Kept as a thin wrapper (rather than calling
// eval.ExactToken directly at each call site) so this file's intent -
// "the traced program's exact output token" - stays named and greppable.
func codeExactAnswer(want string) eval.Evaluator {
	return eval.ExactToken(want)
}

// codeTracePythonSource is the inline Python function traced by
// codeTracePythonTest.
const codeTracePythonSource = `def transform(s, n):
    if n == 0:
        return s
    if len(s) % 2 == 0:
        s = s[1:] + s[0]
    else:
        s = s[::-1]
    return transform(s, n - 1)`

// codeTracePythonWant is the exact result of transform("abcdef", 4).
//
// ground truth: len("abcdef") is 6 (even) at every recursive call, since
// the rotate branch never changes length, so every call takes the rotate
// branch: "abcdef"->"bcdefa"->"cdefab"->"defabc"->"efabcd" over 4 calls
// (n=4,3,2,1), then n=0 returns "efabcd". codebase_trace_test.go both
// hand-derives this independently and, when python3 is available, runs
// the source above directly and compares stdout.
const codeTracePythonWant = "efabcd"

// codeTracePythonTest: trace a recursive Python function combining a
// parity check, string rotation, and full reversal to its exact output.
func codeTracePythonTest() testkit.Test {
	prompt := `Here is a Python function:

` + "```python\n" + codeTracePythonSource + "\n```" + `

What is the exact return value of transform("abcdef", 4)? Respond with
only the resulting string, with no quotes and no other text.`

	return testkit.Test{
		ID:          "code-trace-python",
		Category:    "research",
		Subcategory: "codebase",
		Description: "Trace a recursive Python function combining a parity check, string rotation, and full reversal to its exact output.",
		Prompt:      prompt,
		Eval:        codeExactAnswer(codeTracePythonWant),
	}
}

// codeTraceTSSource is the inline TypeScript function traced by
// codeTraceTSTest.
const codeTraceTSSource = `function pack(items: number[]): string {
  return items.reduce((acc, x, i) => {
    if (i % 2 === 0) return acc + x;
    return acc + x * 2;
  }, 0).toString(2);
}`

// codeTSPackTrace replicates pack's reduce logic in Go so
// codebase_trace_test.go can cross-check the hand-derived trace by
// running equivalent logic, since the TypeScript itself is not executed.
func codeTSPackTrace(items []int) string {
	acc := 0
	for i, x := range items {
		if i%2 == 0 {
			acc += x
		} else {
			acc += x * 2
		}
	}
	return strconv.FormatInt(int64(acc), 2)
}

// codeTraceTSWant is the exact result of pack([3,4,5,6,7]).
//
// ground truth: acc starts at 0; i=0,x=3 (even index) -> acc=3; i=1,x=4
// (odd index) -> acc=3+4*2=11; i=2,x=5 -> acc=11+5=16; i=3,x=6 ->
// acc=16+6*2=28; i=4,x=7 -> acc=28+7=35. 35 in binary is 100011
// (32+2+1=35). codebase_trace_test.go recomputes this both by hand and by
// calling codeTSPackTrace.
const codeTraceTSWant = "100011"

// codeTraceTSTest: trace a TypeScript reduce-based function mixing
// index-parity branching and binary string conversion to its exact
// output.
func codeTraceTSTest() testkit.Test {
	prompt := `Here is a TypeScript function:

` + "```ts\n" + codeTraceTSSource + "\n```" + `

What is the exact return value of pack([3, 4, 5, 6, 7])? Respond with only
the resulting string, with no quotes and no other text.`

	return testkit.Test{
		ID:          "code-trace-ts",
		Category:    "research",
		Subcategory: "codebase",
		Description: "Trace a TypeScript reduce-based function mixing index-parity branching and binary string conversion to its exact output.",
		Prompt:      prompt,
		Eval:        codeExactAnswer(codeTraceTSWant),
	}
}