"""End-to-end tests: run a program directly and through the compiler + VM.

Each case is compiled to a protected chunk, executed with the stock Luau CLI
and compared byte-for-byte with the direct run.
"""
import subprocess
import sys
import tempfile
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

import compiler

LUAU = compiler.LUAU

CASES = {
    'hello': 'print("hello")\n',
    'arith': 'local a,b = 1,2\nprint(a+b, a-b, a*b, a/b, a%b, 2^10, 7//2)\n',
    'function': 'local function f(x) return x*2 end\nprint(f(21))\n',
    'table_method': '''
local t = {1,2,3, x=4}
t.y = t[1] + t.x
function t:m(a) return self.x + a end
print(t:m(5), #t, t.y)
''',
    'numeric_for': '''
local s = 0
for i = 1, 5 do s = s + i end
for i = 10, 1, -2 do s = s + i end
for i = 1.0, 2.5, 0.5 do s = s + i end
print(s)
''',
    'while_repeat': '''
local i, s = 0, 0
while i < 5 do i = i + 1; s = s + i end
repeat s = s - 1 until s < 10
print(i, s)
''',
    'generic_for': '''
local t = {10, 20, 30, a = 1}
local s = 0
for i, v in ipairs(t) do s = s + v end
for k, v in pairs(t) do if type(v) == "number" then s = s + v end end
for i = 1, 3 do s = s + 1 end
print(s)
''',
    'closures': '''
local function counter()
  local n = 0
  return function() n = n + 1 return n end
end
local c = counter()
print(c(), c(), c())
local fs = {}
for i = 1, 3 do fs[i] = function() return i end end
print(fs[1](), fs[2](), fs[3]())
''',
    'upvalue_chain': '''
local x = 1
local function a()
  local y = 2
  local function b()
    local z = 3
    return function() return x + y + z end
  end
  return b()
end
print(a()())
''',
    'varargs': '''
local function f(...)
  local n = select("#", ...)
  local s = 0
  for i = 1, n do s = s + (select(i, ...)) end
  return n, s
end
print(f(1,2,3))
print(f())
local function g(a, ...) return a, ... end
print(g(9, 8, 7))
''',
    'multi_return': '''
local function f() return 1, 2, 3 end
local a, b, c, d = f()
print(a, b, c, d)
local t = {f()}
print(#t, t[1], t[3])
print((f()))
local function g() return f() end
print(g())
''',
    'recursion': '''
local function fib(n) if n < 2 then return n end return fib(n-1) + fib(n-2) end
print(fib(12))
''',
    'strings': '''
local s = "abc" .. "def" .. 1 .. 2
print(s, #s, s:upper(), s:sub(2, 4), string.rep("x", 3))
print(("a" < "b"), ("a" == "a"), tostring(nil), tonumber("42") + 1)
''',
    'logic': '''
local function t(x) return x and "yes" or "no" end
print(t(1), t(nil), t(false), t(""))
print(nil or 5, false and 9, 0 or 7)
''',
    'pcall': '''
local ok, err = pcall(function() error("boom") end)
print(ok, tostring(err):match("boom") ~= nil)
print(pcall(function(a, b) return a + b end, 2, 3))
''',
    'table_lib': '''
local t = {5, 2, 8}
table.insert(t, 1, 9)
table.sort(t)
print(table.concat(t, ","))
local u = {}
u[#u + 1] = "a"
u[#u + 1] = "b"
print(#u, u[1] .. u[2])
''',
    'continue': '''
local s = 0
for i = 1, 10 do
  if i % 2 == 0 then continue end
  s = s + i
end
print(s)
''',
    'nested_calls': '''
local function add3(a, b, c) return a + b + c end
local function id(x) return x end
print(add3(id(1), id(2), add3(1, 2, 3)))
''',
    'setlist_multret': '''
local function f() return 4, 5, 6 end
local t = {1, 2, 3, f()}
print(#t, table.concat(t, "-"))
''',
    'shadow': '''
local x = 1
do
  local x = 2
  print(x)
end
print(x)
''',
    'method_chain': '''
local M = {}
M.__index = M
function M.new(v) return setmetatable({v = v}, M) end
function M:add(n) self.v = self.v + n return self end
function M:get() return self.v end
print(M.new(1):add(2):add(3):get())
''',
    'setlist_gap': '''
local names = {"a", "b", "c"}
local out = {}
for _, n in ipairs(names) do out[n] = function() return n end end
print(out.a(), out.b(), out.c())
''',
    'late_upvalue': '''
local f
local function g() return f() end
f = function() return 7 end
print(g())
''',
    'vararg_forward': '''
local function f(a, b, c, d, e) return a, b, c, d, e end
local function g(...) return f(...) end
print(g(1, 2, 3, 4, 5))
local function wrap(...)
  local result = f(...)
  return result
end
print(wrap(6, 7, 8, 9, 10))
''',
}


def run(path):
    p = subprocess.run([LUAU, str(path)], capture_output=True, text=True)
    return p.returncode, p.stdout, p.stderr


def main(only=None):
    ok = fail = 0
    for name, src in CASES.items():
        if only and name not in only:
            continue
        with tempfile.TemporaryDirectory() as tmp:
            tmp = Path(tmp)
            direct = tmp / 'direct.luau'
            direct.write_text(src)
            rc, out, err = run(direct)
            if rc != 0:
                print('SKIP %-18s (direct run failed: %s)' % (name, err.strip().splitlines()[:1]))
                continue
            prot = tmp / 'prot.luau'
            try:
                prot.write_text(compiler.compile_source(src, key=7, hardened=True, binary=True))
            except Exception as exc:
                print('FAIL %-18s compile: %s' % (name, exc))
                fail += 1
                continue
            rc2, out2, err2 = run(prot)
        if rc2 == 0 and out2 == out:
            print('ok   %-18s %s' % (name, out.strip().replace('\n', ' | ')[:70]))
            ok += 1
        else:
            print('FAIL %-18s rc=%d\n  want: %r\n  got : %r\n  err : %s'
                  % (name, rc2, out, out2, err2.strip().replace('\n', ' ')[:300]))
            fail += 1
    print('\n%d passed, %d failed' % (ok, fail))
    return 0 if fail == 0 else 1


if __name__ == '__main__':
    raise SystemExit(main(sys.argv[1:] or None))