#!/usr/bin/env python3
"""
Tengwar v0.3 Test Suite — All New Features
"""

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from tengwar.interpreter import Interpreter

def run_tests():
    i = Interpreter()
    passed = 0
    failed = 0
    
    def test(category, name, source, expected=None):
        nonlocal passed, failed, i
        try:
            result = i.run_source(source)
            result_str = repr(result)
            if expected is None or result_str == expected:
                passed += 1
                return True
            else:
                print(f'  ✗ [{category}] {name}: expected {expected}, got {result_str}')
                failed += 1
                return False
        except Exception as e:
            print(f'  ✗ [{category}] {name}: {e}')
            failed += 1
            return False

    # ============================================================
    # DICTIONARIES
    # ============================================================
    print("Dictionaries:")
    test("dict", "create", '(dict-size (dict "a" 1 "b" 2))', '2')
    test("dict", "get", '(>> (dict "x" 42) → d (dict-get d "x"))', '42')
    test("dict", "get default", '(dict-get (dict "a" 1) "b" 99)', '99')
    test("dict", "set", '(>> (dict "a" 1) → d (dict-get (dict-set d "b" 2) "b"))', '2')
    test("dict", "del", '(dict-size (dict-del (dict "a" 1 "b" 2) "a"))', '1')
    test("dict", "keys", '(len (dict-keys (dict "a" 1 "b" 2)))', '2')
    test("dict", "vals", '(len (dict-vals (dict "a" 1 "b" 2)))', '2')
    test("dict", "pairs", '(len (dict-pairs (dict "a" 1 "b" 2)))', '2')
    test("dict", "has? true", '(dict-has? (dict "a" 1) "a")', 'true')
    test("dict", "has? false", '(dict-has? (dict "a" 1) "b")', 'false')
    test("dict", "merge", '(dict-size (dict-merge (dict "a" 1) (dict "b" 2)))', '2')
    test("dict", "from vec", '(dict-size (dict ⟦⟨"a" 1⟩ ⟨"b" 2⟩⟧))', '2')
    print(f"  ✓ {passed} dict tests passed")
    dict_count = passed

    # ============================================================
    # LET BINDINGS
    # ============================================================
    print("\nLet Bindings:")
    test("let", "basic", '(let x 10 y 20 (+ x y))', '30')
    test("let", "nested", '(let a 5 (let b 10 (* a b)))', '50')
    test("let", "shadow", '(let x 99 x)', '99')
    test("let", "body expr", '(let n 7 (map {+ _ n} ⟦1 2 3⟧))', '⟦8 9 10⟧')
    let_count = passed - dict_count
    print(f"  ✓ {let_count} let tests passed")

    # ============================================================
    # PIPE OPERATOR
    # ============================================================
    print("\nPipe Operator:")
    test("pipe", "basic", '(pipe 5 {+ _ 1} {* _ 2})', '12')
    test("pipe", "chain", '(pipe ⟦3 1 2⟧ sort head)', '1')
    test("pipe", "filter+map", '(pipe ⟦1 2 3 4 5⟧ {filter {> _ 2} _} {map {* _ 10} _})', '⟦30 40 50⟧')
    pipe_count = passed - dict_count - let_count
    print(f"  ✓ {pipe_count} pipe tests passed")

    # ============================================================
    # ERROR HANDLING
    # ============================================================
    print("\nError Handling:")
    test("errors", "catch success", '(catch (+ 1 2) (λ e e))', '3')
    test("errors", "try success", '(fst (try {+ 1 1} {_}))', 'true')
    err_count = passed - dict_count - let_count - pipe_count
    print(f"  ✓ {err_count} error tests passed")

    # ============================================================
    # TYPE SYSTEM
    # ============================================================
    print("\nType Checking:")
    test("type", "int?", '(int? 42)', 'true')
    test("type", "float?", '(float? 3.14)', 'true')
    test("type", "str?", '(str? "hi")', 'true')
    test("type", "bool?", '(bool? true)', 'true')
    test("type", "vec?", '(vec? ⟦1 2⟧)', 'true')
    test("type", "fn?", '(fn? {+ _ 1})', 'true')
    test("type", "nil?", '(nil? ∅)', 'true')
    test("type", "type int", '(type 42)', '"int"')
    test("type", "type str", '(type "hello")', '"str"')
    test("type", "type vec", '(type ⟦1 2 3⟧)', '"vec"')
    test("type", "type fn", '(type {+ _ 1})', '"fn"')
    test("type", "type dict", '(type (dict "a" 1))', '"dict"')
    type_count = passed - dict_count - let_count - pipe_count - err_count
    print(f"  ✓ {type_count} type tests passed")

    # ============================================================
    # EXTENDED FUNCTIONAL
    # ============================================================
    print("\nExtended Functional:")
    test("func", "flat-map", '(flat-map (λ x ⟦x x⟧) ⟦1 2 3⟧)', '⟦1 1 2 2 3 3⟧')
    test("func", "find", '(find {> _ 3} ⟦1 2 3 4 5⟧)', '4')
    test("func", "find none", '(nil? (find {> _ 99} ⟦1 2 3⟧))', 'true')
    test("func", "index-of", '(index-of {> _ 3} ⟦1 2 3 4 5⟧)', '3')
    test("func", "take", '(take 2 ⟦1 2 3 4 5⟧)', '⟦1 2⟧')
    test("func", "drop", '(drop 3 ⟦1 2 3 4 5⟧)', '⟦4 5⟧')
    test("func", "take-while", '(take-while {< _ 4} ⟦1 2 3 4 5⟧)', '⟦1 2 3⟧')
    test("func", "drop-while", '(drop-while {< _ 4} ⟦1 2 3 4 5⟧)', '⟦4 5⟧')
    test("func", "zip-with", '(zip-with + ⟦1 2 3⟧ ⟦10 20 30⟧)', '⟦11 22 33⟧')
    test("func", "unique", '(unique ⟦1 2 2 3 3 3⟧)', '⟦1 2 3⟧')
    test("func", "partition", '(>> (partition {> _ 3} ⟦1 2 3 4 5⟧) → r (fst r))', '⟦4 5⟧')
    test("func", "scan", '(scan + 0 ⟦1 2 3 4⟧)', '⟦0 1 3 6 10⟧')
    test("func", "chunks", '(chunks 2 ⟦1 2 3 4 5 6⟧)', '⟦⟦1 2⟧ ⟦3 4⟧ ⟦5 6⟧⟧')
    test("func", "interleave", '(interleave ⟦1 2 3⟧ ⟦10 20 30⟧)', '⟦1 10 2 20 3 30⟧')
    test("func", "repeat", '(repeat 3 0)', '⟦0 0 0⟧')
    test("func", "iterate", '(iterate {* _ 2} 1 5)', '⟦1 2 4 8 16⟧')
    test("func", "sort-by", '(sort-by {_ } ⟦3 1 2⟧)', '⟦1 2 3⟧')
    test("func", "for-each", '(nil? (for-each {+ _ 1} ⟦1 2 3⟧))', 'true')
    func_count = passed - dict_count - let_count - pipe_count - err_count - type_count
    print(f"  ✓ {func_count} functional tests passed")

    # ============================================================
    # MATH EXTENDED
    # ============================================================
    print("\nMath Extended:")
    test("math", "abs", '(abs -42)', '42')
    test("math", "abs pos", '(abs 42)', '42')
    test("math", "min", '(min 3 1 2)', '1')
    test("math", "max", '(max 3 1 2)', '3')
    test("math", "min vec", '(min ⟦5 2 8 1⟧)', '1')
    test("math", "max vec", '(max ⟦5 2 8 1⟧)', '8')
    test("math", "clamp low", '(clamp -5 0 10)', '0')
    test("math", "clamp high", '(clamp 15 0 10)', '10')
    test("math", "clamp mid", '(clamp 5 0 10)', '5')
    test("math", "floor", '(floor 3.7)', '3')
    test("math", "ceil", '(ceil 3.2)', '4')
    test("math", "round", '(round 3.5)', '4')
    test("math", "pi exists", '(> pi 3.0)', 'true')
    test("math", "e exists", '(> e 2.0)', 'true')
    math_count = passed - dict_count - let_count - pipe_count - err_count - type_count - func_count
    print(f"  ✓ {math_count} math tests passed")

    # ============================================================
    # STRING EXTENDED
    # ============================================================
    print("\nString Extended:")
    test("str", "starts-with?", '(starts-with? "hello world" "hello")', 'true')
    test("str", "ends-with?", '(ends-with? "hello.tw" ".tw")', 'true')
    test("str", "chars", '(chars "abc")', '⟦"a" "b" "c"⟧')
    test("str", "char-at", '(char-at "hello" 1)', '"e"')
    test("str", "fmt", '(fmt "{}+{}={}" 1 2 3)', '"1+2=3"')
    test("str", "pad-left", '(pad-left "42" 5 "0")', '"00042"')
    test("str", "pad-right", '(pad-right "hi" 5 ".")', '"hi..."')
    str_count = passed - dict_count - let_count - pipe_count - err_count - type_count - func_count - math_count
    print(f"  ✓ {str_count} string tests passed")

    # ============================================================
    # JSON
    # ============================================================
    print("\nJSON:")
    test("json", "encode vec", '(json-encode ⟦1 2 3⟧)', '"[1, 2, 3]"')
    test("json", "encode str", '(str? (json-encode "hello"))', 'true')
    test("json", "parse obj", '(type (json-parse "{}"))', '"dict"')
    test("json", "parse arr", '(type (json-parse "[]"))', '"vec"')
    test("json", "roundtrip", '(json-encode (json-parse "[1,2,3]"))', '"[1, 2, 3]"')
    json_count = passed - dict_count - let_count - pipe_count - err_count - type_count - func_count - math_count - str_count
    print(f"  ✓ {json_count} JSON tests passed")

    # ============================================================
    # SYSTEM
    # ============================================================
    print("\nSystem:")
    test("sys", "time-ms", '(> (time-ms) 0)', 'true')
    test("sys", "uuid length", '(> (len (uuid)) 30)', 'true')
    test("sys", "rand range", '(>> (rand) → r (& (>= r 0.0) (< r 1.0)))', 'true')
    sys_count = passed - dict_count - let_count - pipe_count - err_count - type_count - func_count - math_count - str_count - json_count
    print(f"  ✓ {sys_count} system tests passed")

    # ============================================================
    # PYTHON INTEROP
    # ============================================================
    print("\nPython Interop:")
    test("py", "import math", '(>> (py-import "math") (py-call math "sqrt" 16.0))', '4.0')
    test("py", "py-eval", '(py-eval "2 ** 10")', '1024')
    test("py", "py-attr", '(> (py-attr (py-import "math") "pi") 3.0)', 'true')
    test("py", "dot access", '(>> (py-import "math") (> math.pi 3.0))', 'true')
    py_count = passed - dict_count - let_count - pipe_count - err_count - type_count - func_count - math_count - str_count - json_count - sys_count
    print(f"  ✓ {py_count} python interop tests passed")

    # ============================================================
    # DEEP RECURSION
    # ============================================================
    print("\nDeep Recursion:")
    i_rec = Interpreter()
    i_rec.run_source('(↺ cd (λ n (? (= n 0) "done" (cd (- n 1)))))')
    r = i_rec.run_source('(cd 10000)')
    if repr(r) == '"done"':
        passed += 1
        print("  ✓ 10000 recursive calls")
    else:
        failed += 1
        print(f"  ✗ 10000 recursive calls: {repr(r)}")

    i_rec.run_source('(↺ sum-to (λ n acc (? (= n 0) acc (sum-to (- n 1) (+ acc n)))))')
    r2 = i_rec.run_source('(sum-to 5000 0)')
    expected_sum = sum(range(5001))
    if repr(r2) == str(expected_sum):
        passed += 1
        print(f"  ✓ sum 1..5000 = {expected_sum}")
    else:
        failed += 1
        print(f"  ✗ sum 1..5000: expected {expected_sum}, got {repr(r2)}")

    # ============================================================
    # SANDBOX
    # ============================================================
    print("\nSandbox Mode:")
    sandbox = Interpreter(sandbox=True)
    
    for name, code in [
        ("blocks file I/O", '(read-file "test.txt")'),
        ("blocks network", '(http-get "http://example.com")'),
        ("blocks py-import", '(py-import "os")'),
        ("blocks py-eval", '(py-eval "import os")'),
    ]:
        try:
            sandbox.run_source(code)
            print(f"  ✗ sandbox {name}")
            failed += 1
        except Exception:
            print(f"  ✓ sandbox {name}")
            passed += 1

    sandbox2 = Interpreter(sandbox=True)
    sandbox2.max_steps = 200
    try:
        sandbox2.run_source('(↺ loop (λ n (loop (+ n 1))))')
        sandbox2.run_source('(loop 0)')
        print("  ✗ step limit")
        failed += 1
    except Exception:
        print("  ✓ step limit enforced")
        passed += 1

    # ============================================================
    # BINARY AST PROTOCOL
    # ============================================================
    print("\nBinary AST Protocol:")
    from tengwar.binary_ast import encode, decode, run_binary, encode_b64, run_b64, ASTBuilder

    # Roundtrip tests
    for name, src, expected in [
        ("simple add", "(+ 1 2)", "3"),
        ("string", '(concat "hello" " world")', '"hello world"'),
        ("vector", "(map {* _ 2} ⟦1 2 3⟧)", "⟦2 4 6⟧"),
        ("let binding", "(let x 10 y 20 (+ x y))", "30"),
        ("pipe", "(pipe ⟦3 1 2⟧ sort head)", "1"),
        ("conditional", "(? (> 5 3) \"yes\" \"no\")", '"yes"'),
        ("nested", "(reduce + 0 (filter {> _ 2} ⟦1 2 3 4 5⟧))", "12"),
    ]:
        binary = encode(src)
        result = run_binary(binary)
        if repr(result) == expected:
            passed += 1
            text_bytes = len(src.encode('utf-8'))
            print(f"  ✓ {name} [{len(binary)}B binary vs {text_bytes}B text]")
        else:
            failed += 1
            print(f"  ✗ {name}: expected {expected}, got {repr(result)}")

    # Base64 roundtrip
    b64 = encode_b64("(+ 40 2)")
    r = run_b64(b64)
    if repr(r) == "42":
        passed += 1
        print(f"  ✓ base64 roundtrip")
    else:
        failed += 1
        print(f"  ✗ base64 roundtrip: {repr(r)}")

    # ASTBuilder
    b = ASTBuilder()
    prog = b.program(
        b.define("square", b.lam(["x"], b.binop("*", b.sym("x"), b.sym("x")))),
        b.apply(b.sym("square"), b.int(7))
    )
    r = b.run(prog)
    if repr(r) == "49":
        passed += 1
        print("  ✓ ASTBuilder: square(7) = 49")
    else:
        failed += 1
        print(f"  ✗ ASTBuilder: {repr(r)}")

    # Builder → binary → run
    binary = b.encode(prog)
    r = run_binary(binary)
    if repr(r) == "49":
        passed += 1
        print(f"  ✓ Builder→binary→run [{len(binary)}B]")
    else:
        failed += 1

    # ============================================================
    # CONVERSION HELPERS
    # ============================================================
    print("\nConversion Helpers:")
    test("conv", "dict-to-vec", '(len (dict-to-vec (dict "a" 1 "b" 2)))', '2')
    test("conv", "vec-to-dict", '(dict-get (vec-to-dict ⟦⟨"x" 42⟩⟧) "x")', '42')

    # ============================================================
    # SUMMARY
    # ============================================================
    print(f"\n{'='*50}")
    print(f"v0.3 RESULTS: {passed} passed, {failed} failed, {passed + failed} total")
    if failed == 0:
        print("ALL v0.3 TESTS PASSED! ✓")
    else:
        print(f"FAILURES: {failed}")
    print(f"{'='*50}")
    return failed == 0

# (run_tests called from bottom of file)


# ==========================================
# v0.3.1 — TOKEN OPTIMIZATION TESTS
# ==========================================

def run_v031_tests():
    """Tests for v0.3.1 optimizations: flat cond, builtins, VM"""
    from tengwar.interpreter import Interpreter

    i = Interpreter()

    tests = [
        # --- Flat cond ---
        ("flat cond basic", '(cond (> 5 3) "yes" _ "no")', '"yes"'),
        ("flat cond multi", '(cond (= 1 2) "a" (= 2 2) "b" _ "c")', '"b"'),
        ("flat cond FizzBuzz 15", '(let n 15 (cond (divides? n 15) "FB" (divides? n 3) "F" (divides? n 5) "B" _ n))', '"FB"'),
        ("flat cond FizzBuzz 9", '(let n 9 (cond (divides? n 15) "FB" (divides? n 3) "F" (divides? n 5) "B" _ n))', '"F"'),
        ("flat cond FizzBuzz 10", '(let n 10 (cond (divides? n 15) "FB" (divides? n 3) "F" (divides? n 5) "B" _ n))', '"B"'),
        ("flat cond FizzBuzz 7", '(let n 7 (cond (divides? n 15) "FB" (divides? n 3) "F" (divides? n 5) "B" _ n))', '7'),
        ("flat cond default only", '(cond _ 42)', '42'),

        # --- Numeric predicates ---
        ("zero? true", "(zero? 0)", "true"),
        ("zero? false", "(zero? 5)", "false"),
        ("pos? true", "(pos? 5)", "true"),
        ("pos? false", "(pos? -3)", "false"),
        ("neg? true", "(neg? -3)", "true"),
        ("neg? false", "(neg? 5)", "false"),
        ("even? true", "(even? 4)", "true"),
        ("even? false", "(even? 3)", "false"),
        ("odd? true", "(odd? 3)", "true"),
        ("odd? false", "(odd? 4)", "false"),
        ("divides? true", "(divides? 15 3)", "true"),
        ("divides? false", "(divides? 15 4)", "false"),
        ("empty? vec true", "(empty? ⟦⟧)", "true"),
        ("empty? vec false", "(empty? ⟦1⟧)", "false"),
        ("empty? str true", '(empty? "")', "true"),
        ("empty? str false", '(empty? "a")', "false"),

        # --- Compact math builtins ---
        ("inc", "(inc 41)", "42"),
        ("dec", "(dec 43)", "42"),
        ("sqr", "(sqr 7)", "49"),
        ("sqr float", "(sqr 2.5)", "6.25"),
        ("cube", "(cube 3)", "27"),
        ("pow", "(pow 2 10)", "1024"),
        ("sqrt", "(sqrt 49.0)", "7.0"),
        ("double", "(double 21)", "42"),
        ("half", "(half 10)", "5.0"),
        ("neg", "(neg 5)", "-5"),
        ("id", "(id 42)", "42"),

        # --- sqr in pipelines ---
        ("sum of squares", "(sum (map sqr (range 1 11)))", "385"),
        ("map double", "(map double ⟦1 2 3⟧)", "⟦2 4 6⟧"),
        ("filter+inc", "(map inc (filter even? ⟦1 2 3 4 5⟧))", "⟦3 5⟧"),

        # --- Memo ---
        ("memo fib 10", "(>> (≡ fib (memo (λ n (? (< n 2) n (+ (fib (- n 1)) (fib (- n 2))))))) (fib 10))", "55"),
        ("memo fib 20", "(>> (≡ fib (memo (λ n (? (< n 2) n (+ (fib (- n 1)) (fib (- n 2))))))) (fib 20))", "6765"),

        # --- Multi-expression ---
        ("multi-expr defn", "(defn sq x (* x x)) (sq 7)", "49"),
        ("multi-expr chain", "(≡ a 10) (≡ b 20) (+ a b)", "30"),

        # --- for comprehension ---
        ("for basic", "(sum (for x (range 1 6) (* x x)))", "55"),
        ("for+when", "(sum (for x (range 1 11) when (even? x) (* x 2)))", "60"),
        ("for+when+sqr", "(sum (for x (range 1 11) when (odd? x) (sqr x)))", "165"),
    ]

    passed = 0
    failed = 0
    for name, code, expected in tests:
        try:
            result = repr(i.run_source(code))
            if result == expected:
                passed += 1
            else:
                print(f"  ✗ {name}: got {result}, expected {expected}")
                failed += 1
        except Exception as e:
            print(f"  ✗ {name}: ERROR {e}")
            failed += 1

    # VM tests
    from tengwar.vm import VM
    vm = VM(Interpreter())
    vm_tests = [
        ("vm int", "42", "42"),
        ("vm add", "(+ 1 2)", "3"),
        ("vm nested math", "(+ (* 3 4) (- 10 5))", "17"),
        ("vm cond true", "(? (> 5 3) 10 20)", "10"),
        ("vm cond false", "(? (< 5 3) 10 20)", "20"),
        ("vm vector", "⟦1 2 3⟧", "⟦1 2 3⟧"),
        ("vm string", '"hello"', '"hello"'),
        ("vm let", "(let x 10 (+ x 5))", "15"),
        ("vm define+use", "(>> (≡ x 42) x)", "42"),
        ("vm builtin", "(len ⟦1 2 3⟧)", "3"),
        ("vm sum", "(sum ⟦1 2 3⟧)", "6"),
        ("vm map", "(map inc ⟦1 2 3⟧)", "⟦2 3 4⟧"),
    ]
    for name, code, expected in vm_tests:
        try:
            result = repr(vm.run_source(code))
            if result == expected:
                passed += 1
            else:
                print(f"  ✗ {name}: got {result}, expected {expected}")
                failed += 1
        except Exception as e:
            print(f"  ✗ {name}: ERROR {e}")
            failed += 1

    total = passed + failed
    print(f"\nv0.3.1 RESULTS: {passed} passed, {failed} failed, {total} total")
    if failed == 0:
        print(f"ALL v0.3.1 TESTS PASSED! ✓")
    return failed == 0




if __name__ == "__main__":
    s1 = run_tests()
    print()
    s2 = run_v031_tests()
    sys.exit(0 if (s1 and s2) else 1)
