"""Future execution guard: exact inequality certificates; unresolved grades raise.

This does not modify the active EI reference wrapper. Per-grade timeout is enforced
by a separate process boundary, including native C operations. No timeout becomes
a reward bit, and no failed worker is automatically restarted.
"""
import argparse
import json
import os
from pathlib import Path
import runpy
import select
import subprocess
import sys
import time

from certified_integer_grade import unequal_certificate
from run_with_certified_grader import UnresolvedGrade


class IsolatedGrader:
    def __init__(self, normalize, log_path, seconds=5):
        self.normalize = normalize
        self.log_path = Path(log_path)
        self.log_path.parent.mkdir(parents=True, exist_ok=True)
        self.seconds = seconds
        assert seconds > 0
        self.sequence = 0
        self.closed = False
        self.stderr = self.log_path.with_suffix('.worker.stderr.log').open('ab')
        env = os.environ.copy()
        env.update(CUDA_VISIBLE_DEVICES='', OMP_NUM_THREADS='1', MKL_NUM_THREADS='1')
        self.process = subprocess.Popen([sys.executable, str(Path(__file__).with_name('isolated_grade_worker.py'))],
                                        env=env, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                                        stderr=self.stderr, bufsize=0)
        try:
            ready = self._receive(30)
            assert ready['status'] == 'ready'
            self.log(dict(status='worker_ready', **{k:v for k,v in ready.items() if k!='status'},
                          grade_timeout_seconds=seconds))
        except BaseException:
            self.close()
            raise

    def log(self, record):
        with self.log_path.open('a') as stream:
            stream.write(json.dumps(dict(time=time.time(), **record))+'\n')

    def _receive(self, timeout):
        deadline = time.monotonic() + timeout
        received = bytearray()
        while b'\n' not in received:
            remaining = deadline - time.monotonic()
            if remaining <= 0 or not select.select([self.process.stdout], [], [], remaining)[0]:
                raise UnresolvedGrade('Isolated symbolic grading exceeded the wall-clock deadline')
            chunk = os.read(self.process.stdout.fileno(), 4096)
            if not chunk:
                raise UnresolvedGrade('Isolated symbolic grading worker exited without a grade')
            received.extend(chunk)
            if len(received) > 65536:
                raise UnresolvedGrade('Isolated grading response exceeded message bound')
        line, extra = received.split(b'\n', 1)
        if extra:
            raise UnresolvedGrade('Unexpected extra worker response')
        return json.loads(line)

    def grade(self, prediction, reference):
        if self.closed:
            raise UnresolvedGrade('Isolated grading service is closed; no automatic restart')
        certificate = unequal_certificate(self.normalize(prediction), self.normalize(reference))
        if certificate is not None:
            self.log(dict(status='certified_unequal', prediction=prediction, reference=reference,
                          certificate=certificate))
            return False
        self.sequence += 1
        started = time.perf_counter()
        try:
            self.process.stdin.write((json.dumps(dict(id=self.sequence, prediction=prediction, reference=reference))+'\n').encode())
            self.process.stdin.flush()
            reply = self._receive(self.seconds)
            if reply.get('id') != self.sequence or reply.get('status') != 'graded' or type(reply.get('result')) is not bool:
                raise UnresolvedGrade('Isolated grader returned no grade: '+json.dumps(reply))
            return reply['result']
        except BaseException as error:
            self.log(dict(status='unresolved', prediction=prediction, reference=reference,
                          seconds=time.perf_counter()-started, error=repr(error)))
            self.close()
            if isinstance(error, (KeyboardInterrupt, SystemExit)):
                raise
            raise UnresolvedGrade('No reward emitted after isolated grading failure') from error

    def close(self):
        if self.closed:
            return
        self.closed = True
        if self.process.poll() is None:
            self.process.terminate()
            try:
                self.process.wait(timeout=2)
            except subprocess.TimeoutExpired:
                self.process.kill()
                self.process.wait(timeout=2)
        for stream in (self.process.stdin, self.process.stdout, self.stderr):
            try:
                stream.close()
            except (BrokenPipeError, OSError):
                pass


def main():
    started = time.perf_counter()
    parser = argparse.ArgumentParser()
    parser.add_argument('--certificate-log', type=Path, required=True)
    parser.add_argument('--timing-receipt', type=Path)
    parser.add_argument('script', type=Path)
    parser.add_argument('arguments', nargs=argparse.REMAINDER)
    args = parser.parse_args()
    if args.timing_receipt:
        assert not args.timing_receipt.exists(), 'Never overwrite prior execution accounting'
    service = None
    completed = False
    try:
        import run_baseline as base
        from reasoning_from_scratch import ch03
        service = IsolatedGrader(ch03.normalize_text, args.certificate_log)
        ch03.grade_answer = service.grade
        base.grade_answer = service.grade
        sys.argv = [str(args.script), *args.arguments]
        runpy.run_path(str(args.script), run_name='__main__')
        completed = True
    finally:
        if service:
            service.close()
        if args.timing_receipt:
            with args.timing_receipt.open('x') as stream:
                json.dump(dict(status='complete' if completed else 'failed', script=str(args.script),
                               wrapper_wall_seconds=time.perf_counter()-started), stream, indent=2)


if __name__ == '__main__':
    main()
