#!/usr/bin/python3
# Simple asciicast v2 encoder (see https://asciinema.org/)
# Author: David Coles <coles.david@gmail.com>

import argparse
import curses
import json
import locale
import os
import time

STDIN_FILENO = 0
STDERR_FILENO = 2
READBUF = 4096


def main():
    curses.setupterm()

    parser = argparse.ArgumentParser()
    parser.add_argument(
            '-W', '--width', type=int, default=curses.tigetnum('cols'))
    parser.add_argument(
            '-H', '--height', type=int, default=curses.tigetnum('lines'))
    parser.add_argument('-S', '--shell', default=os.environ['SHELL'])
    parser.add_argument('-T', '--term', default=os.environ['TERM'])
    parser.add_argument('-r', '--force-cr', action='store_true',
            help='Force newlines to be treated as CR+NL')
    parser.add_argument('-e', '--echo', action='store_true',
            help='Echo input to stderr')
    parser.add_argument('--encoding', default=locale.getpreferredencoding())
    args = parser.parse_args()

    # Use raw STDIN/STDERR to avoid buffering and control encoding
    stdin = os.fdopen(STDIN_FILENO, 'rb', buffering=0)
    stderr = os.fdopen(STDERR_FILENO, 'wb')

    start_time = time.time()

    # Output in asciicast v2 format:
    # https://github.com/asciinema/asciinema/blob/develop/doc/asciicast-v2.md
    print(json.dumps({
        'version': 2,
        'width': args.width,
        'height': args.height,
        'timestamp': int(start_time),
        'env': {
            'SHELL': args.shell,
            'TERM': args.term,
        },
    }))

    for buf in iter(lambda: stdin.read(READBUF), b''):
        ts = time.time() - start_time

        if args.echo:
            stderr.write(buf)
            stderr.flush()

        if args.force_cr:
            buf = buf.replace(b'\n', b'\r\n')

        print(json.dumps([round(ts, 6), 'o', buf.decode(args.encoding)]),
              flush=True)


if __name__ == '__main__':
    main()
