#!/usr/bin/env python3
"""array receiver calls: execution, type checks and mutation contracts."""
import argparse
import re
import subprocess

parser = argparse.ArgumentParser()
parser.add_argument('host', nargs='?', default='bin/zy')
parser.add_argument('--case', default='')
args = parser.parse_args()

cases = [
    ('growable', True, 'include rt; fn main() { int[?] a = {2,3}; a.push(7); rt.put_n(a.len()); rt.put_c(32); rt.put_n(a.cap() >= a.len()); rt.put_c(32); rt.put_n(a.pop()); a.resize(4); rt.put_c(32); rt.put_n(a[3]); rt.put_c(32); rt.put_n(a.len()); }', '3 1 7 0 4'),
    ('fixed length', True, 'include rt; fn main() { int[3] a = {2,3,4}; rt.put_n(a.len()); rt.put_c(32); rt.put_n(a.cap()); fixed int[?] b = rt.new(fixed int[?], 5); rt.put_c(32); rt.put_n(b.len()); }', '3 3 5'),
    ('packed', True, 'include rt; fn main() { u8[?] a = {1}; a.push(300); rt.put_n(a.pop()); rt.put_c(32); i16[?] b = {-5}; rt.put_n(b.pop()); rt.put_c(32); f32[?] c = {2.5}; c.push(1.5); rt.put_n(c.pop() == 1.5); }', '44 -5 1'),
    ('optimized', True, 'include rt; o3 int change(int[?] a) { a.push(9); int n = a.pop(); a.resize(3); return n + a.len(); } fn main() { int[?] a = {1}; rt.put_n(change(a)); rt.put_c(32); rt.put_n(a.len()); }', '12 3'),
    ('field receiver', True, 'include rt; struct box { int[?] items; }; fn main() { box b; b.items = rt.new(int[?]); b.items.push(8); rt.put_n(b.items.pop()); }', '8'),
    ('nested arguments', True, 'include rt; fn main() { int[?] a = {1}, b = {2,3}; a.push(b.pop() + b.len()); a.resize(b.len() + 2); rt.put_n(a[1]); rt.put_c(32); rt.put_n(a.len()); }', '4 3'),
    ('fixed push', False, 'fn main() { int[2] a; a.push(1); }', None),
    ('fixed pop', False, 'fn main() { int[2] a; return a.pop(); }', None),
    ('fixed resize', False, 'fn main() { int[2] a; a.resize(3); }', None),
    ('fixed handle push', False, 'include rt; fn main() { fixed int[?] a = rt.new(fixed int[?], 2); a.push(1); }', None),
    ('element mismatch', False, 'include rt; fn main() { int[?] a = {}; a.push("wrong"); }', None),
    ('extra argument', False, 'include rt; fn main() { int[?] a = {}; a.len(1); }', None),
    ('invalidated borrow', False, 'include rt; fn main() { int[?] a = {1}; int* p = &a[0]; a.resize(10); return *p; }', None),
    ('free fixed push', False, 'fn main() { int[2] a; push(a,1); }', None),
    ('free fixed pop', False, 'fn main() { int[2] a; return pop(a); }', None),
    ('registered generic', True, 'include rt, run; T first<T>(fixed T[?] a) { return a[0]; } immediate fn register() { run.next(); run.add_method("first", "first"); } register fn main() { int[?] a = {9}; rt.put_n(a.first()); int[2] b = {7,8}; rt.put_c(32); rt.put_n(b.first()); }', '9 7'),
    ('registered callback', True, 'include rt, run; U apply<T,U>(fixed T[?] a, U(T) f) { return f(a[0]); } immediate fn register() { run.next(); run.add_method("apply", "apply"); } register fn main() { int[?] a = {4}; rt.put_n(a.apply(fn (e) { return e + 2; })); rt.put_c(32); float result = a.apply(fn:float(int e) { return e * 0.5; }); rt.put_n(result == 2.0); }', '6 1'),
    ('registered callback mismatch', False, 'include rt, run; U apply<T,U>(fixed T[?] a, U(T) f) { return f(a[0]); } immediate fn register() { run.next(); run.add_method("apply", "apply"); } register fn main() { int[?] a = {4}; return a.apply(fn:int(char* e) { return e[0]; }); }', None),
    ('registration requires receiver', False, 'include run; int empty() { return 0; } immediate fn register() { run.next(); run.add_method("empty", "empty"); } register fn main() { return 0; }', None),
    ('declarative registration', True, 'include rt, run; T first<T>(fixed T[?] a) { return a[0]; } run.method("first", first); fn main() { int[3] a = {6,7,8}; rt.put_n(a.first()); }', '6'),
    ('map inferred', True, 'include rt, std; fn main() { int[?] a = {1,2,3}; int[?] b = a.map(fn(e) { return e + 2; }); rt.put_n(b[0]); rt.put_c(32); rt.put_n(b[2]); rt.put_c(32); float[?] c = a.map(fn(e) { return e * 0.5; }); rt.put_n(c[2] == 1.5); }', '3 5 1'),
    ('filter fixed', True, 'include rt, std; fn main() { int[5] a = {1,2,3,4,5}; let b = a.filter(fn(e) { return e % 2; }); rt.put_n(b.len()); rt.put_c(32); rt.put_n(b[2]); }', '3 5'),
    ('search methods', True, 'include rt, std; fn main() { int[?] a = {2,4,6}; rt.put_n(a.find(4)); rt.put_c(32); rt.put_n(a.find(5)); rt.put_c(32); rt.put_n(a.contains(6)); rt.put_c(32); rt.put_n(a.find_if(fn(e) { return e > 4; })); }', '1 -1 1 2'),
    ('reduce and predicates', True, 'include rt, std; fn main() { int[3] a = {2,4,6}; rt.put_n(a.reduce(1, fn(acc,e) { return acc + e; })); rt.put_c(32); rt.put_n(a.all(fn(e) { return e % 2 == 0; })); rt.put_c(32); rt.put_n(a.any(fn(e) { return e > 5; })); }', '13 1 1'),
    ('empty algorithms', True, 'include rt, std; fn main() { int[?] a = {}; int calls = 0; let b = a.map(fn(e) { calls++; return e; }); let c = a.filter(fn(e) { calls++; return e; }); rt.put_n(b.len() + c.len() + calls); rt.put_c(32); rt.put_n(a.all(fn(e) { return e; })); rt.put_c(32); rt.put_n(a.any(fn(e) { return e; })); rt.put_c(32); rt.put_n(a.reduce(9, fn(acc,e) { return acc + e; })); }', '0 1 0 9'),
    ('shared filter references', True, 'include rt, std; fn main() { char*[?] a = rt.new(char*[?]); char* s = rt.new(char, 2); s[0] = 65; a.push(s); let b = a.filter(fn(e) { return e[0] == 65; }); rt.free(a); rt.gc(); rt.put_c(b[0][0]); }', 'A'),
    ('map needs import', False, 'include rt; fn main() { int[?] a = {1}; let b = a.map(fn(e) { return e; }); }', None),
    ('pop transfers ownership', True, 'include rt; fn main() { int before = rt.live_objects(); { int*[?] a = rt.new(int*[?]); int* p = rt.new(int); *p = 7; a.push(p); let saved = a.pop(); rt.free(a); rt.put_n(rt.isobj(saved)); rt.put_c(32); rt.put_n(*saved); } rt.put_c(32); rt.put_n(rt.live_objects() - before); }', '1 7 0'),
    ('returned pop survives array', True, 'include rt; int*.u take() { int*[?] a = rt.new(int*[?]); int* p = rt.new(int); *p = 8; a.push(p); return a.pop(); } fn main() { let p = take(); rt.put_n(rt.isobj(p)); rt.put_c(32); rt.put_n(*p); }', '1 8'),
    ('single field struct push', False, 'include rt; struct item { int value; }; fn main() { item[?] a = rt.new(item[?]); item x; x.value = 7; a.push(x); }', None),
    ('single field struct pop', False, 'include rt; struct item { int value; }; fn main() { item[?] a = rt.new(item[?], 1); return a.pop().value; }', None),
    ('resize drops children', True, 'include rt; fn main() { int before = rt.live_objects(); int*[?] a = rt.new(int*[?]); a.push(rt.new(int)); a.resize(0); rt.put_n(rt.live_objects() - before); }', '2'),
    ('callback borrows references', False, 'include rt, std; fn main() { int*[?] a = rt.new(int*[?]); a.push(rt.new(int)); let b = a.filter(fn(e) { rt.free(e); return 1; }); }', None),
    ('packed map and filter', True, 'include rt, std; fn main() { i16[?] a = {-5,7}; i16[?] b = a.filter(fn(e) { return e < 0; }); rt.put_n(b[0]); rt.put_c(32); f32[2] c = {1.5,2.5}; f32[?] d = c.map(fn(e) { return e; }); rt.put_n(d[1] == 2.5); }', '-5 1'),
    ('numeric coercion', True, 'include rt; fn main() { float[?] a = rt.new(float[?]); a.push(2); rt.put_n(a[0] == 2.0); rt.put_c(32); int[?] b = {}; b.push(2.5); b.resize(3.0); rt.put_n(b[0]); rt.put_c(32); rt.put_n(b.len()); }', '1 2 3'),
    ('invalid resize length', False, 'include rt; fn main() { int[?] a = {}; a.resize("bad"); }', None),
    ('free pop ownership', True, 'include rt; fn main() { int*[?] a = rt.new(int*[?]); int* p = rt.new(int); *p = 9; push(a,p); let saved = pop(a); rt.free(a); rt.put_n(rt.isobj(saved)); rt.put_c(32); rt.put_n(*saved); }', '1 9'),
    ('packed string values', True, 'include rt; struct box { bstr text; }; fn main() { bstr[?] a = {b"hi"}; let s = a.pop(); rt.put_n(s[1]); rt.put_c(32); box b; b.text = b"ok"; rt.put_n(b.text[1]); }', '105 107'),
    ('field array borrow', False, 'include rt; struct box { int[?] items; }; fn main() { box* b = rt.new(box); b.items = rt.new(int[?], 1); let borrowed = b.items; rt.free(b); return borrowed[0]; }', None),
]

count = 0
for name, accept, source, expected in cases:
    if args.case and args.case not in name:
        continue
    source = 'typecheck on;\n' + source
    check = subprocess.run([args.host], input=source, text=True, capture_output=True, timeout=10)
    diagnostics = check.stdout + check.stderr
    failed = check.returncode != 0 or bool(re.search(r'(zc:|zc0:|error:)', diagnostics))
    assert failed != accept, f'{name}: expected {"accept" if accept else "reject"}\n{diagnostics}'
    if accept:
        assert check.returncode == 0 and check.stdout == expected, f'{name}: {check.stdout!r}, {check.stderr!r}; expected {expected!r}'
    print('ok   array methods: ' + name)
    count += 1
assert count, 'no cases selected'