"""Offline, bounded OCR. Run only inside the restricted native OCR sandbox."""
import base64
import io
import json
import os
import pathlib
import resource
import signal
import subprocess
import sys
import tempfile
import time
import warnings

from PIL import Image

MAX_INPUT = 10 * 1024 * 1024
MAX_TEXT = 50000
MAX_PIXELS = 16_000_000
Image.MAX_IMAGE_PIXELS = MAX_PIXELS
warnings.simplefilter('error', Image.DecompressionBombWarning)


def limits():
    resource.setrlimit(resource.RLIMIT_CORE, (0, 0))
    resource.setrlimit(resource.RLIMIT_FSIZE, (12 * 1024 * 1024,) * 2)
    resource.setrlimit(resource.RLIMIT_NOFILE, (64, 64))
    resource.setrlimit(resource.RLIMIT_AS, (512 * 1024 * 1024,) * 2)
    resource.setrlimit(resource.RLIMIT_CPU, (70, 70))
    resource.setrlimit(resource.RLIMIT_NPROC, (32, 32))
    if hasattr(os, 'sched_getaffinity'):
        os.sched_setaffinity(0, {min(os.sched_getaffinity(0))})
    os.umask(0o077)


def extract(data, mime):
    if not 0 < len(data) <= MAX_INPUT or mime not in ('application/pdf', 'image/png', 'image/jpeg'):
        raise ValueError('Invalid input')
    deadline = time.monotonic() + 75

    def tool(argv, output, maximum, seconds):
        timeout = min(seconds, deadline - time.monotonic())
        if timeout <= 0:
            raise ValueError('Deadline')
        with output.open('wb') as stream:
            subprocess.run(argv, stdin=subprocess.DEVNULL, stdout=stream,
                           stderr=subprocess.DEVNULL, timeout=timeout, check=True,
                           env={'PATH': '/usr/bin:/bin', 'LANG': 'C.UTF-8', 'OMP_THREAD_LIMIT': '1', 'HOME': '/tmp'})
        if output.stat().st_size > maximum:
            raise ValueError('Output limit')
        return output.read_bytes()

    with tempfile.TemporaryDirectory(prefix='ocr-') as directory:
        root = pathlib.Path(directory)
        source = root / 'input.pdf'
        if mime == 'application/pdf':
            if not data.startswith(b'%PDF-'):
                raise ValueError('PDF signature')
            source.write_bytes(data)
            info = tool(['pdfinfo', str(source)], root / 'info', 64000, 5).decode('utf-8', errors='strict')
            pages = next(int(line.split(':', 1)[1].strip()) for line in info.splitlines() if line.startswith('Pages:'))
            if not 1 <= pages <= 10:
                raise ValueError('Page limit')
        else:
            pages = 1
        texts = []
        length = 0
        for page in range(1, pages + 1):
            image = root / 'page.png'
            if mime == 'application/pdf':
                tool(['pdftoppm', '-f', str(page), '-l', str(page), '-singlefile', '-scale-to', '2000', '-png', str(source), str(root / 'page')], root / 'render-log', 0, 10)
                raw = image.read_bytes()
            else:
                raw = data
            # Validate dimensions/format before decoding and strip metadata by re-encoding.
            with Image.open(io.BytesIO(raw)) as decoded:
                expected = ('PNG',) if mime in ('application/pdf', 'image/png') else ('JPEG',)
                if decoded.format not in expected or decoded.width * decoded.height > MAX_PIXELS or max(decoded.size) > 8000 or min(decoded.size) < 1:
                    raise ValueError('Image limits')
                decoded.convert('RGB').save(image, format='PNG')
            text = tool(['tesseract', str(image), 'stdout', '-l', 'eng', '--psm', '3'], root / 'text', MAX_TEXT, 12).decode('utf-8', errors='strict')
            length += len(text.encode('utf-8')) + 1
            if length > MAX_TEXT:
                raise ValueError('Text limit')
            texts.append(text)
            image.unlink()
        return '\n'.join(texts)


def main():
    limits()
    signal.signal(signal.SIGALRM, lambda *_: os._exit(124))
    signal.alarm(85)
    try:
        if os.geteuid() == 0:
            raise ValueError('Run OCR without root')
        if len(sys.argv) == 2 and sys.argv[1] == '--self-test':
            import socket
            root_readonly = False
            try:
                pathlib.Path('/ocr-probe').write_text('x')
            except OSError:
                root_readonly = True
            isolated = not pathlib.Path('/home').exists() and not pathlib.Path('/var/www').exists()
            network_blocked = False
            try:
                connection = socket.create_connection(('1.1.1.1', 443), timeout=1)
                connection.close()
            except OSError:
                network_blocked = True
            subprocess.run(['tesseract', '--version'], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True, timeout=3)
            subprocess.run(['pdftoppm', '-v'], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True, timeout=3)
            if not root_readonly or not isolated or not network_blocked:
                raise ValueError('Sandbox unavailable')
            print(json.dumps({'ready': True, 'uid': os.geteuid(), 'root_readonly': root_readonly, 'network_blocked': network_blocked, 'private_filesystem': isolated, 'memory_limit': resource.getrlimit(resource.RLIMIT_AS)[0], 'process_limit': resource.getrlimit(resource.RLIMIT_NPROC)[0]}))
            return 0
        raw = sys.stdin.buffer.read(14 * 1024 * 1024 + 1)
        if len(raw) > 14 * 1024 * 1024:
            raise ValueError('Request limit')
        request = json.loads(raw)
        data = base64.b64decode(request['data'], validate=True)
        result = {'text': extract(data, request['mime'])}
        print(json.dumps(result, ensure_ascii=True))
    except Exception:
        print('{"error":"OCR could not safely process this file"}')
        return 1
    return 0


if __name__ == '__main__':
    sys.exit(main())