#!/usr/bin/env python3
"""Compare helper calls, shared references, and exceptional heaps with Python."""
import ast
import json
from pathlib import Path
import runpy
import subprocess
import sys
import tempfile

from test_support import project_temporary_directory

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT))
from lang.py.object_snapshot import snapshot_roots
from lang.py.support import SupportChecker

api = runpy.run_path(str(ROOT / 'formalsw'))


def constant(value):
    if value is None:
        return '.none'
    if type(value) is bool:
        return '.bool ' + str(value).lower()
    if type(value) is int:
        return f'.int ({value})'
    return '.string ' + json.dumps(value, ensure_ascii=False)


def data(value):
    if type(value) is list:
        return '.list [' + ', '.join(map(data, value)) + ']'
    return '.scalar (' + constant(value) + ')'


source = '''
def push(items: list[int], value: int = 10) -> list[int]:
    items.append(value)
    return items

def touch(items: list[int], value: int) -> int:
    items.append(value)
    return value

def fail(items: list[int], value: int) -> int:
    items.append(value)
    return value // 0

def recursive(items: list[int], n: int) -> list[int]:
    if n <= 0:
        return items
    push(items, n)
    return recursive(items, n - 1)
'''
bodies = {
    'defaults': 'return push(items)',
    'returned_alias': 'other = push(items, n)\nother.append(n + 1)\nreturn alias',
    'argument_order': 'return touch(items, n) + touch(alias, n + 1)',
    'rhs_mutation': 'items.append(touch(items, n))\nreturn items',
    'argument_failure': 'items.append(fail(alias, n))\nreturn items',
    'receiver_failure': 'n.append(fail(items, n))\nreturn items',
    'left_error': 'return fail(items, n) + touch(alias, n)',
    'right_error': 'return touch(items, n) + fail(alias, n)',
    'new_list': 'return push([n], n + 1)',
    'nested_aliases': 'return [push(items, n), push(alias, n + 1)]',
    'cycle': 'items.append(items)\nreturn items',
    'recursion': 'return recursive(items, n)',
    'inplace': 'items += push(alias, n)\nreturn items',
    'repeat': 'items *= n\nreturn alias',
    'repeat_bool': 'items *= False\nreturn alias',
    'extend_string': 'items += "λ😀"\nreturn items',
    'read_after': 'push(items, n)\nreturn alias[-1]',
    'list_truth': 'if items:\n    push(alias, n)\nreturn items',
    'loop': 'while n > 0:\n    push(items, n)\n    n -= 1\nreturn alias',
    'loop_else': 'while n > 0:\n    push(items, n)\n    break\nelse:\n    push(alias, 0)\nreturn items',
    'conditional': 'return push(items, n) if n > 0 else push(alias)',
    'short_circuit': 'return items and push(alias, n)',
    'index_side_effect': 'return items[touch(alias, n)]',
    'range_loop': 'for k in range(n):\n    push(items, k)\nreturn alias',
    'range_negative': 'for k in range(n, -3, -2):\n    push(items, k)\nreturn items',
    'range_boolean': 'for k in range(False, True):\n    push(items, k)\nreturn items',
    'range_arguments': 'for k in range(touch(items, n), touch(alias, n + 2)):\n    push(items, k)\nreturn items',
    'range_argument_failure': 'for k in range(touch(items, n), fail(alias, n)):\n    push(items, k)\nreturn items',
    'range_zero_step': 'for k in range(touch(items, n), touch(alias, n + 2), 0):\n    push(items, k)\nreturn items',
    'range_invalid_index': 'for k in range(items, touch(alias, n)):\n    push(items, k)\nreturn items',
    'range_else': 'for k in range(n):\n    push(items, k)\nelse:\n    push(alias, 7)\nreturn items',
    'range_break': 'for k in range(n):\n    push(items, k)\n    break\nelse:\n    push(alias, 7)\nreturn items',
    'range_continue': 'for k in range(n):\n    push(items, k)\n    continue\n    push(alias, 99)\nelse:\n    push(alias, 7)\nreturn items',
    'range_nested': 'for k in range(n):\n    for j in range(k, n):\n        push(items, k * 10 + j)\nreturn alias',
    'range_rebind_target': 'for k in range(n):\n    push(items, k)\n    k = 42\nreturn items',
    'range_else_break_outer': 'for k in range(n):\n    for j in range(0):\n        push(items, j)\n    else:\n        push(items, k)\n        break\nelse:\n    push(alias, 7)\nreturn items',
    'reversed_copy': 'other = list(reversed(push(items, n)))\nother.append(8)\nreturn [items, other]',
    'reversed_string': 'return list(reversed("λ😀"))',
    'reversed_type_error': 'return list(reversed(touch(items, n)))',
    'reversed_aliases': 'return list(reversed([items, alias]))',
    'equal_lists': 'return items == alias',
    'unequal_lists': 'return items != alias',
    'equal_numbers': 'return [True, False, n] == [1, 0, n]',
    'equal_mixed': 'return [n, "a", None] == [n, "a", None]',
    'equal_list_scalar': 'return items == n',
    'unequal_scalar_list': 'return n != items',
    'equal_nested_aliases': 'return [items] == [items]',
    'equal_cycle': 'items.append(items)\nreturn items == items',
    'equal_length_mismatch': 'return [[n]] == [[n], 1]',
    'equal_order': 'return push(items, n) == push(alias, n + 1)',
    'maximum_values': 'return max(n, 3, -2)',
    'maximum_order': 'return max(touch(items, n), touch(alias, n + 1))',
    'maximum_bool_tie': 'return max(True, 1)',
    'maximum_int_tie': 'return max(1, True)',
    'maximum_strings': 'return max("α", "😀", "a")',
    'maximum_error': 'return max(touch(items, n), "a", touch(alias, n + 1))',
}
for name, body in bodies.items():
    source += f'\ndef {name}(n: int, items: list[int], alias: list[int]):\n'
    source += ''.join('    ' + line + '\n' for line in body.splitlines())
program = ast.parse(source)
checker = SupportChecker(program)
for name in bodies:
    issues = checker.check(name, heap_calls=True)
    assert not issues, (name, issues)
range_source = 'def f(n: int) -> list[int]:\n    items = []\n    for k in range(n):\n        items.append(k)\n    return items\n'
assert not SupportChecker(ast.parse(range_source)).check('f')
for suffix in ('range = f\n', 'from other import range\n', 'from other import *\n'):
    assert SupportChecker(ast.parse(range_source + suffix)).check('f')
for iterator in ('range()', 'range(n, stop=0)', 'range(n, n, n, n)', '[n]'):
    assert SupportChecker(ast.parse(range_source.replace('range(n)', iterator))).check('f')
reverse_source = 'def f(n: int) -> list[int]: return list(reversed([n]))\n'
assert not SupportChecker(ast.parse(reverse_source)).check('f')
for builtin in ('list', 'reversed'):
    for suffix in (f'{builtin} = f\n', f'from other import {builtin}\n'):
        module = ast.parse(reverse_source + suffix)
        assert SupportChecker(module).check('f')
        assert api['AstJsonEmitter'](module).expr(module.body[0].body[0].value)['type'] == 'Call'
    assert SupportChecker(ast.parse(reverse_source.replace('n: int', f'n: int, {builtin}: int'))).check('f')
maximum_source = 'def f(n: int) -> int: return max(n, 1)\n'
assert not SupportChecker(ast.parse(maximum_source)).check('f')
for expression in ('max(n)', 'max(n, default=0)', 'max(n, 1, key=n)'):
    assert SupportChecker(ast.parse(maximum_source.replace('max(n, 1)', expression))).check('f')
for suffix in ('max = f\n', 'from other import max\n', 'from other import *\n'):
    assert SupportChecker(ast.parse(maximum_source + suffix)).check('f')
maximum_module = ast.parse(maximum_source + 'def g(n: int) -> int: return f(n)\n')
for name in ('f', 'g'):
    issues = SupportChecker(maximum_module).requirement(name,
        f'theorem wrong : Not (scalarFunction Probe.module "{name}" [.int 0]).Halts')
    assert any('positional max requires' in issue.reason for issue in issues)
custom_max = ast.parse('def max(n: int) -> int: return n\ndef f(n: int) -> int: return max(n)\n')
assert not SupportChecker(custom_max).check('f')
assert api['AstJsonEmitter'](custom_max).expr(custom_max.body[1].body[0].value)['type'] == 'Call'
lines = ['import FormalSW', 'open FormalSW.PyAst', 'set_option maxRecDepth 8192',
         'def source : Module := pyast_from_json% (include_str "source.json")']
count = 0
for name in bodies:
    for initial in ([], [4]):
        for n in (-1, 0, 2):
            for shared in (False, True):
                world = {}
                exec(source, world)
                anchors = [initial.copy(), [9]]
                try:
                    value = world[name](n, anchors[0], anchors[0 if shared else 1])
                    error = None
                except (IndexError, TypeError, AttributeError, ZeroDivisionError, ValueError) as exception:
                    value, error = None, type(exception).__name__
                expected = data(snapshot_roots(anchors + ([value] if error is None else [])))
                heap = '[' + ', '.join('[' + ', '.join('.scalar (' + constant(v) + ')' for v in xs) + ']'
                                      for xs in (initial, [9])) + ']'
                call = f'HeapCalls.call source {json.dumps(name)} [.scalar (.int ({n})), .ref 0, .ref {0 if shared else 1}] {heap}'
                outcome = ('match c.outcome with | .ok _ => true | _ => false' if error is None else
                           f'match c.outcome with | .error e => decide (e = ⟨{json.dumps(error)}, []⟩) | _ => false')
                lines.extend([
                    f'-- {name}, {initial!r}, n={n}, shared={shared}',
                    f'example : (match ({call}).execution.observe 1000 with',
                    f'  | .done (.ok c) => ({outcome}) &&',
                    f'      (Graph.observedBool ((Graph.compare c.snapshot ({expected})).observe 5000) == some true)',
                    '  | _ => false) = true := by native_decide',
                ])
                count += 1
with tempfile.TemporaryDirectory(prefix='formalsw-heap-calls-') as directory:
    directory = Path(directory)
    (directory / 'source.json').write_text(api['AstJsonEmitter'](program).emit())
    (directory / 'Checks.lean').write_text('\n'.join(lines) + '\n')
    subprocess.run(['lake', 'env', 'lean', '-DwarningAsError=true', str(directory / 'Checks.lean')], cwd=ROOT, check=True)

# Partial heap operations must not turn an unsupported computation into a proof
# of divergence or equivalence. Coverage is independent of Not and Requires.
with project_temporary_directory(prefix='heap-coverage-') as directory:
    directory = Path(directory)
    for index, marker in enumerate(('Returns', 'Not Halts', 'Equiv(f)', 'Not Equiv(f)')):
        path = directory / f'probe{index}.py'
        path.write_text(f'def f(n: int) -> bool:\n    """\n    @formalsw: {marker}\n    """\n    return [n] < [n]\n')
        module = api['source_module'](path)
        obligations = api['collect_requirements'](path, module, 'Probe', 'ProbeAst')
        coverage = [item for item in obligations if item.marker == 'execution model coverage']
        assert len(coverage) == 1 and '.Supported' in coverage[0].property_type
        assert 'Not' not in coverage[0].property_type
    path = directory / 'vacuous.py'
    path.write_text('def f(n: int) -> bool:\n    """\n    @formalsw: Equiv(f)\n    """\n    return [n] < [n]\n')
    stub = subprocess.run([str(ROOT / 'formalsw'), '--stub', str(path)], cwd=ROOT,
                          text=True, capture_output=True, check=True).stdout
    before = stub.split('theorem f.«Supported»')[0]
    namespace = next(line.split(' ', 1)[1] for line in stub.splitlines() if line.startswith('namespace '))
    path.with_suffix('.py.lean').write_text(before.replace('  sorry', '  intro n outcome\n  rfl') + f'end {namespace}\n')
    rejected = subprocess.run([str(ROOT / 'formalsw'), str(path)], cwd=ROOT, text=True, capture_output=True)
    assert rejected.returncode != 0 and 'Supported' in rejected.stdout + rejected.stderr
    assert 'Unknown identifier' in rejected.stdout + rejected.stderr, rejected.stdout + rejected.stderr

    # Existing interpreters cannot be used to bypass the coverage obligation.
    module = ast.parse('def f(n: int) -> list[int]: return [n]\n')
    issues = SupportChecker(module).requirement('f', 'theorem wrong : Not (HeapCalls.scalarFunction Probe.module "f" []).Halts')
    assert any('list results require HeapCalls.call' in issue.reason for issue in issues)
    for model in ('Calls.function', 'scalarFunction', 'Mutable.function'):
        issues = SupportChecker(module).requirement('f', f'theorem wrong : Not ({model} Probe.module \"f\" []).Halts')
        assert any('require the heap call interpreter' in issue.reason for issue in issues), model
    globals_module = ast.parse('items: list[int] = []\ndef f(n: int) -> None: items.append(n)\n')
    issues = SupportChecker(globals_module).requirement('f', 'theorem wrong : Not (HeapCalls.call Probe.module "f" [] []).Halts')
    assert any('does not support globals' in issue.reason for issue in issues)

subprocess.run([sys.executable, str(ROOT / 'tests/py/heap_calls.py')], check=True)
print(f'Heap helper calls: {count} Python comparisons, aliases, exception ordering, and coverage checks passed')
