package eval

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"
	"sync"
	"time"
)

// leadingPackageClausePattern matches a Go source line that is nothing but
// a package clause (optionally indented), used to defensively strip a
// model's own "package xxx" line before splicing its code into
// solution.go. Only the first non-blank line is ever checked/stripped; a
// package clause appearing later (inside a string, comment, or as a
// genuine syntax error) is left alone.
var leadingPackageClausePattern = regexp.MustCompile(`^\s*package\s+\w+\s*$`)

// stripLeadingPackageClause removes a leading "package xxx" line from code,
// if the first non-blank line is one. This lets a model's response be
// either a bare code fragment or a full file (package clause + imports +
// declarations) and still compile as part of solution.go's package main.
func stripLeadingPackageClause(code string) string {
	lines := strings.Split(code, "\n")
	for i, line := range lines {
		if strings.TrimSpace(line) == "" {
			continue
		}
		if leadingPackageClausePattern.MatchString(line) {
			return strings.Join(append(lines[:i:i], lines[i+1:]...), "\n")
		}
		break
	}
	return code
}

// execTimeout bounds every run step (and the C compile) of the exec
// evaluators; goBuildTimeout is the longer allowance for a cold `go build`.
const (
	execTimeout    = 30 * time.Second
	goBuildTimeout = 120 * time.Second
)

// execWaitDelay bounds how long Wait lingers for output pipes after the
// process is killed; without it a grandchild holding stdout open hangs the
// whole run past execTimeout.
const execWaitDelay = 2 * time.Second

// maxExecOutput caps captured stdout/stderr per call, so a runaway print
// loop cannot exhaust memory before the timeout fires.
const maxExecOutput = 1 << 20

// execSlots limits concurrent compile/run subprocesses. The runner issues up
// to Concurrency scoring calls at once; unbounded, CPU starvation pushes
// correct solutions past execTimeout and turns load into false zeros.
var execSlots = make(chan struct{}, max(2, runtime.NumCPU()/2))

// cappedBuffer keeps the first maxExecOutput bytes written and discards the
// rest, still reporting full writes so the child never sees a broken pipe.
type cappedBuffer struct{ buf bytes.Buffer }

func (c *cappedBuffer) Write(p []byte) (int, error) {
	if room := maxExecOutput - c.buf.Len(); room > 0 {
		c.buf.Write(p[:min(len(p), room)])
	}
	return len(p), nil
}

func (c *cappedBuffer) String() string { return c.buf.String() }

// runBounded runs name with args in dir under env, holding an exec slot,
// with a hard timeout and capped output. A timeout is reported as such
// rather than as a bare "signal: killed".
func runBounded(ctx context.Context, timeout time.Duration, dir string, env []string, name string, args ...string) (stdout, stderr string, err error) {
	select {
	case execSlots <- struct{}{}:
		defer func() { <-execSlots }()
	case <-ctx.Done():
		return "", "", ctx.Err()
	}

	runCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	// #nosec G204 -- running model-generated code under an isolated temp dir
	// and a minimal env allowlist is this package's purpose; name is always
	// a fixed toolchain binary or one this package just compiled. See the
	// README "Requirements" warning (S13).
	cmd := exec.CommandContext(runCtx, name, args...) //nolint:gosec // see comment above
	cmd.Dir = dir
	cmd.Env = env
	cmd.WaitDelay = execWaitDelay
	var out, errOut cappedBuffer
	cmd.Stdout = &out
	cmd.Stderr = &errOut
	err = cmd.Run()
	if err != nil && runCtx.Err() == context.DeadlineExceeded {
		err = fmt.Errorf("timed out after %s: %w", timeout, err)
	}
	return out.String(), errOut.String(), err
}

// CodePlaceholder marks where the extracted response code is substituted
// into a harness template.
const CodePlaceholder = "{{CODE}}"

// PassthroughHarness is a harness that runs the extracted code verbatim,
// with no surrounding driver. Use it when the response is expected to be a
// complete, self-contained program (e.g. a full Python script) rather than
// a fragment plugged into a test driver.
const PassthroughHarness = CodePlaceholder

// buildSource substitutes code into harness at CodePlaceholder.
func buildSource(harness, code string) string {
	return strings.ReplaceAll(harness, CodePlaceholder, code)
}

// goEnvCache memoizes `go env GOCACHE`/`go env GOMODCACHE` so every
// exec-evaluator call does not pay for an extra subprocess: they are read
// once per process and reused.
var (
	goEnvOnce     sync.Once
	goCacheDir    string
	goModCacheDir string
)

// goEnvValue shells out to `go env key` and returns its trimmed output, or
// "" if the go tool is unavailable or the call fails.
func goEnvValue(key string) string {
	// #nosec G204 -- key is always one of the two fixed, hardcoded literals
	// ("GOCACHE"/"GOMODCACHE") passed by minimalExecEnv below, never
	// response/user-controlled input.
	out, err := exec.Command("go", "env", key).Output() //nolint:gosec // key is one of two hardcoded literals, see comment above
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(out))
}

// minimalExecEnv builds a minimal environment allowlist for compiling and
// running model-generated code, rather than inheriting the operator's full
// environment: PATH (needed to find go/python3/cc and any tool they shell
// out to), HOME and TMPDIR pinned to homeDir (an isolated, single-call temp
// directory, never the operator's real home), and GOCACHE/GOMODCACHE
// inherited from the operator's actual Go environment so `go run` reuses
// the existing build/module cache instead of rebuilding it from scratch on
// every call.
func minimalExecEnv(homeDir string) []string {
	goEnvOnce.Do(func() {
		goCacheDir = goEnvValue("GOCACHE")
		goModCacheDir = goEnvValue("GOMODCACHE")
	})

	// Pinned so output does not vary with the operator's shell: timezone,
	// hash seeding and locale-dependent encoding. GOTOOLCHAIN=local stops
	// `go` fetching a different toolchain over the network mid-run.
	env := []string{
		"PATH=" + os.Getenv("PATH"),
		"HOME=" + homeDir,
		"TMPDIR=" + homeDir,
		"TZ=UTC",
		"PYTHONHASHSEED=0",
		"PYTHONUTF8=1",
		"PYTHONDONTWRITEBYTECODE=1",
		"GOTOOLCHAIN=local",
		"GOFLAGS=-buildvcs=false",
	}
	if goCacheDir != "" {
		env = append(env, "GOCACHE="+goCacheDir)
	}
	if goModCacheDir != "" {
		env = append(env, "GOMODCACHE="+goModCacheDir)
	}
	return env
}

// mkWorkDirs creates an isolated per-call temp root plus a "work"
// subdirectory inside it. Callers write source files into work and run
// commands with cmd.Dir = work, while pointing HOME/TMPDIR (via
// minimalExecEnv) at root, never at work itself: the Go toolchain
// specifically refuses to detect a go.mod when the module directory is
// exactly equal to the resolved system temp root ("ignoring go.mod in
// system temp root"), which fires precisely when TMPDIR is set to the same
// directory `go run` is invoked from. Keeping root and work as distinct
// (parent, child) directories avoids that while still confining
// HOME/TMPDIR to a directory nothing else uses.
func mkWorkDirs(prefix string) (root, work string, err error) {
	root, err = os.MkdirTemp("", prefix)
	if err != nil {
		return "", "", err
	}
	work = filepath.Join(root, "work")
	if err := os.Mkdir(work, 0o700); err != nil {
		return "", "", err
	}
	return root, work, nil
}

// checkToolchain reports a Skipped score if name is not on PATH, nil
// otherwise.
func checkToolchain(name string) *Score {
	if _, err := exec.LookPath(name); err != nil {
		s := Score{Skipped: true, Detail: fmt.Sprintf("toolchain missing: %s", name)}
		return &s
	}
	return nil
}

// compareStdout scores full credit when stdout, trimmed, equals want,
// trimmed. cmdErr, when non-nil, always scores zero and surfaces stderr.
func compareStdout(stdout, want string, cmdErr error, stderr string) Score {
	if cmdErr != nil {
		return Score{Value: 0, Detail: fmt.Sprintf("execution failed: %v; stderr: %s", cmdErr, strings.TrimSpace(stderr))}
	}
	got := strings.TrimSpace(stdout)
	wantTrimmed := strings.TrimSpace(want)
	if got == wantTrimmed {
		return Score{Value: 1, Detail: "stdout matches expected"}
	}
	return Score{Value: 0, Detail: fmt.Sprintf("stdout = %q, want %q", got, wantTrimmed)}
}

// goRunEval compiles/runs a two-file Go module: solution.go (the model's
// code, wrapped in package main) plus harness.go (a complete, independent
// driver file supplied by the test author).
type goRunEval struct {
	harness string
	want    string
}

// GoRun returns an Evaluator that writes the response's first fenced Go
// code block to its own solution.go (defensively stripping a leading
// "package xxx" line and prefixing "package main"), writes harness
// verbatim to its own harness.go (a complete file: its own package clause,
// imports, and func main), `go run`s both together in an isolated temp
// module, and compares trimmed stdout to want. Splitting the model's code
// and the harness into separate files means each may import what it needs
// independently, instead of splicing code into a single shared import
// block. If the `go` toolchain is unavailable, the Score is Skipped rather
// than zero.
func GoRun(harness, want string) Evaluator {
	return goRunEval{harness: harness, want: want}
}

func (g goRunEval) Evaluate(ctx context.Context, response string) Score {
	if s := checkToolchain("go"); s != nil {
		return *s
	}

	code := ExtractCodeBlock(response, "go")
	solution := "package main\n\n" + stripLeadingPackageClause(code)

	root, work, err := mkWorkDirs("llmtb-go-*")
	if err != nil {
		return Score{Value: 0, Detail: fmt.Sprintf("mkdtemp: %v", err)}
	}
	defer func() { _ = os.RemoveAll(root) }() // best-effort cleanup; failure is not actionable

	if writeErr := os.WriteFile(filepath.Join(work, "solution.go"), []byte(solution), 0o600); writeErr != nil {
		return Score{Value: 0, Detail: fmt.Sprintf("write solution.go: %v", writeErr)}
	}
	if writeErr := os.WriteFile(filepath.Join(work, "harness.go"), []byte(g.harness), 0o600); writeErr != nil {
		return Score{Value: 0, Detail: fmt.Sprintf("write harness.go: %v", writeErr)}
	}
	if modErr := os.WriteFile(filepath.Join(work, "go.mod"), []byte("module llmtbharness\n\ngo 1.23\n"), 0o600); modErr != nil {
		return Score{Value: 0, Detail: fmt.Sprintf("write go.mod: %v", modErr)}
	}

	// Build and run as separate steps so a cold compile under load is not
	// charged against the run timeout.
	env := minimalExecEnv(root)
	binPath := filepath.Join(work, "prog")
	if _, buildStderr, buildErr := runBounded(ctx, goBuildTimeout, work, env, "go", "build", "-o", binPath, "."); buildErr != nil {
		return Score{Value: 0, Detail: fmt.Sprintf("go build failed: %v; stderr: %s", buildErr, strings.TrimSpace(buildStderr))}
	}
	stdout, stderr, err := runBounded(ctx, execTimeout, work, env, binPath)

	return compareStdout(stdout, g.want, err, stderr)
}

// pyRunEval runs a Python program built from a harness template plus the
// extracted Python code block from the response.
type pyRunEval struct {
	harness string
	want    string
}

// PyRun returns an Evaluator that substitutes the response's first fenced
// Python code block into harness at CodePlaceholder, runs it with python3,
// and compares trimmed stdout to want. If python3 is unavailable, the Score
// is Skipped rather than zero.
func PyRun(harness, want string) Evaluator {
	return pyRunEval{harness: harness, want: want}
}

func (p pyRunEval) Evaluate(ctx context.Context, response string) Score {
	if s := checkToolchain("python3"); s != nil {
		return *s
	}

	code := ExtractCodeBlock(response, "python")
	source := buildSource(p.harness, code)

	root, work, err := mkWorkDirs("llmtb-py-*")
	if err != nil {
		return Score{Value: 0, Detail: fmt.Sprintf("mkdtemp: %v", err)}
	}
	defer func() { _ = os.RemoveAll(root) }() // best-effort cleanup; failure is not actionable

	scriptPath := filepath.Join(work, "script.py")
	if writeErr := os.WriteFile(scriptPath, []byte(source), 0o600); writeErr != nil {
		return Score{Value: 0, Detail: fmt.Sprintf("write script.py: %v", writeErr)}
	}

	stdout, stderr, err := runBounded(ctx, execTimeout, work, minimalExecEnv(root), "python3", scriptPath)

	return compareStdout(stdout, p.want, err, stderr)
}

// cRunEval compiles/runs a C program built from a harness template plus the
// extracted C code block from the response.
type cRunEval struct {
	harness string
	want    string
}

// CRun returns an Evaluator that substitutes the response's first fenced C
// code block into harness at CodePlaceholder, compiles it with cc, runs the
// binary, and compares trimmed stdout to want. If cc is unavailable, the
// Score is Skipped rather than zero.
func CRun(harness, want string) Evaluator {
	return cRunEval{harness: harness, want: want}
}

func (c cRunEval) Evaluate(ctx context.Context, response string) Score {
	if s := checkToolchain("cc"); s != nil {
		return *s
	}

	code := ExtractCodeBlock(response, "c")
	source := buildSource(c.harness, code)

	root, work, err := mkWorkDirs("llmtb-c-*")
	if err != nil {
		return Score{Value: 0, Detail: fmt.Sprintf("mkdtemp: %v", err)}
	}
	defer func() { _ = os.RemoveAll(root) }() // best-effort cleanup; failure is not actionable

	srcPath := filepath.Join(work, "main.c")
	if writeErr := os.WriteFile(srcPath, []byte(source), 0o600); writeErr != nil {
		return Score{Value: 0, Detail: fmt.Sprintf("write main.c: %v", writeErr)}
	}
	binPath := filepath.Join(work, "prog")

	env := minimalExecEnv(root)
	if _, buildStderr, buildErr := runBounded(ctx, execTimeout, work, env, "cc", "-o", binPath, srcPath); buildErr != nil {
		return Score{Value: 0, Detail: fmt.Sprintf("cc failed: %v; stderr: %s", buildErr, strings.TrimSpace(buildStderr))}
	}

	stdout, stderr, err := runBounded(ctx, execTimeout, work, env, binPath)

	return compareStdout(stdout, c.want, err, stderr)
}