"""the corpus + the objective oracle.

each task is a real one-shot zig coding ask plus a hidden `test` (the oracle) and a
known-good `reference` (used only to validate the oracle). tasks are intentionally
reference-dependent: langref-documented features a weak model tends to misremember.
oracle inputs are forced to runtime (`var x = ...; _ = &x;`) so `zig test` truly
*executes* the function rather than comptime-folding a literal.
"""

from __future__ import annotations

import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path


@dataclass(frozen=True)
class Task:
    id: str
    prompt: str
    oracle: str  # a `test {...}` referencing the required decl
    reference: str  # a correct solution; oracle-validation only
    note: str = ""


TASKS: list[Task] = [
    Task("intcast",
         "Implement `pub fn toU8(x: u32) u8` that narrows x to a u8 using @intCast (assume x fits).",
         'test "o" { const s=@import("std"); var x: u32 = 200; _ = &x; try s.testing.expectEqual(@as(u8, 200), toU8(x)); }',
         "pub fn toU8(x: u32) u8 { return @intCast(x); }",
         "@intCast single-arg form (old: @intCast(u8, x))"),
    Task("for-index",
         "Implement `pub fn indexOf(haystack: []const u8, needle: u8) ?usize` returning the index of the first occurrence of needle, or null.",
         'test "o" { const s=@import("std"); var buf=[_]u8{\'a\',\'b\',\'c\'}; const h: []const u8 = &buf; var n: u8=\'c\'; _=&n; try s.testing.expectEqual(@as(?usize,2), indexOf(h,n)); var m: u8=\'z\'; _=&m; try s.testing.expectEqual(@as(?usize,null), indexOf(h,m)); }',
         "pub fn indexOf(haystack: []const u8, needle: u8) ?usize { for (haystack, 0..) |c, i| { if (c == needle) return i; } return null; }",
         "for-loop index capture `for (xs, 0..) |x, i|`"),
    Task("switch-enum",
         "Define `const Dir = enum { north, south, east, west };` and `pub fn dx(d: Dir) i8` returning 1 for east, -1 for west, 0 otherwise, using a switch.",
         'test "o" { const s=@import("std"); var d: Dir = .east; _=&d; try s.testing.expectEqual(@as(i8,1), dx(d)); d=.west; try s.testing.expectEqual(@as(i8,-1), dx(d)); d=.north; try s.testing.expectEqual(@as(i8,0), dx(d)); }',
         "const Dir = enum { north, south, east, west }; pub fn dx(d: Dir) i8 { return switch (d) { .east => 1, .west => -1, .north, .south => 0 }; }",
         "exhaustive switch on enum"),
    Task("labeled-block",
         "Implement `pub fn maxOf(a: i32, b: i32) i32` returning the larger value, computed as a labeled block expression (`blk: { ... break :blk ...; }`).",
         'test "o" { const s=@import("std"); var a: i32=3; var b: i32=7; _=&a; _=&b; try s.testing.expectEqual(@as(i32,7), maxOf(a,b)); a=9; b=2; try s.testing.expectEqual(@as(i32,9), maxOf(a,b)); }',
         "pub fn maxOf(a: i32, b: i32) i32 { return blk: { if (a > b) break :blk a; break :blk b; }; }",
         "labeled block + break :label value"),
    Task("error-union",
         "Define `const E = error{ Bad };` and `pub fn addOk(a: E!i32, b: E!i32) E!i32` that returns the sum, propagating either error with `try`.",
         'test "o" { const s=@import("std"); var a: i32=2; var b: i32=3; _=&a; _=&b; try s.testing.expectEqual(@as(i32,5), try addOk(a,b)); try s.testing.expectError(E.Bad, addOk(E.Bad,b)); }',
         "const E = error{ Bad }; pub fn addOk(a: E!i32, b: E!i32) E!i32 { return (try a) + (try b); }",
         "error union + try propagation"),
    Task("sentinel-slice",
         "Implement `pub fn slice3(arr: *const [3:0]u8) []const u8` returning the 3 elements as a slice (excluding the sentinel).",
         'test "o" { const s=@import("std"); var a: [3:0]u8 = .{ \'a\', \'b\', \'c\' }; _=&a; try s.testing.expectEqualSlices(u8, "abc", slice3(&a)); }',
         "pub fn slice3(arr: *const [3:0]u8) []const u8 { return arr[0..3]; }",
         "sentinel-terminated array, slicing"),
    Task("cstrlen",
         "Implement `pub fn cstrLen(s: [*:0]const u8) usize` returning the number of bytes before the 0 sentinel (like C strlen).",
         'test "o" { const s=@import("std"); var buf=[_:0]u8{\'h\',\'e\',\'l\',\'l\',\'o\'}; _=&buf; const p: [*:0]const u8 = &buf; try s.testing.expectEqual(@as(usize,5), cstrLen(p)); }',
         "pub fn cstrLen(s: [*:0]const u8) usize { var i: usize = 0; while (s[i] != 0) i += 1; return i; }",
         "sentinel-terminated many-item pointer"),
]


def zig_test_source(source: str, zig: str = "zig", timeout: int = 60) -> tuple[bool, str]:
    """compile and run `zig test` on a source string; (passed, last_diagnostic_line)."""
    with tempfile.TemporaryDirectory() as d:
        f = Path(d) / "t.zig"
        f.write_text(source)
        try:
            r = subprocess.run([zig, "test", str(f)], capture_output=True, text=True, timeout=timeout)
        except FileNotFoundError:
            raise FileNotFoundError(f"zig not found ('{zig}')")
        except subprocess.TimeoutExpired:
            return False, "zig test timed out"
    out = (r.stderr or r.stdout).strip()
    return r.returncode == 0, (out.splitlines() or [""])[-1]


def grade(code: str, task: Task, zig: str = "zig") -> tuple[bool, str]:
    """run the task's hidden oracle against agent-produced code."""
    return zig_test_source(code + "\n" + task.oracle, zig=zig)


def validate_oracles(zig: str = "zig") -> list[tuple[Task, bool, str]]:
    """every oracle must accept its reference solution (run before spending tokens)."""
    return [(t, *zig_test_source(t.reference + "\n" + t.oracle, zig=zig)) for t in TASKS]
