#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys
import os
import time
import subprocess
import logging
import argparse
import json
import datetime
import io

logging.basicConfig(format='%(message)s', level=logging.DEBUG)
logger = logging.getLogger(__name__)

def main():
    parser = argparse.ArgumentParser(description='Capture screenshot and videos on Steam Frame device')
    parser.add_argument('--filename', '-f',
                       default='/tmp/screenshot.png',
                       help='Output')
    parser.add_argument('--timestamp', action='store_true',
                       help='Add timestamp')
    parser.add_argument('--json', action='store_true',
                       help='Output result as JSON')

    args = parser.parse_args()

    output_buffer = io.StringIO()
    try:
        steamvr_path = subprocess.check_output(['steamvr', 'path'], stderr=subprocess.STDOUT, universal_newlines=True).strip()
        cdd = os.path.join(steamvr_path, 'bin/linuxarm64')
        run_vrcmd = os.path.join(cdd, 'vrcmd')
        assert os.path.exists(run_vrcmd), "vrcmd not found"

        # Enable recording
        cmd = [run_vrcmd, '--mailboxcmd', 'vrcompositor_systemlayer', 'set_local_video_record?enabled=true']
        output_buffer.write(f"Command: {' '.join(cmd)}\n")
        result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
        output_buffer.write(result.stdout)
        if result.returncode != 0:
            raise subprocess.CalledProcessError(result.returncode, cmd)

        # Wait for the video device to produce frames.
        # Note that even when disabled it outputs roughly 2 blank frames per second.
        cmd = ['timeout', '1', 'ffmpeg', '-f', 'v4l2', '-i', '/dev/video99', '-frames:v', '4', '-f', 'null', '-', '-v', 'error']
        output_buffer.write(f"Command: {' '.join(cmd)}\n")
        retries = 2
        while True:
            result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
            output_buffer.write(result.stdout)
            if result.returncode == 0:
                break
            retries -= 1
            if retries <= 0:
                raise Exception("Failed to get video frames from /dev/video99. Is VR active? Is the v4l2 configuration correct?")

        output_filename = args.filename
        if args.timestamp:
            timestamp = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
            base, ext = os.path.splitext(output_filename)
            output_filename = f"{base}-{timestamp}{ext}"

        # Capture screenshot
        cmd = ['ffmpeg', '-f', 'v4l2', '-i', '/dev/video99', '-frames:v', '1', '-q:v', '1', '-y', output_filename]
        output_buffer.write(f"Command: {' '.join(cmd)}\n")
        result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
        output_buffer.write(result.stdout)
        if result.returncode != 0:
            raise subprocess.CalledProcessError(result.returncode, cmd)

        # Disable recording
        cmd = [run_vrcmd, '--mailboxcmd', 'vrcompositor_systemlayer', 'set_local_video_record?enabled=false']
        output_buffer.write(f"Command: {' '.join(cmd)}\n")
        result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
        output_buffer.write(result.stdout)
        if result.returncode != 0:
            raise subprocess.CalledProcessError(result.returncode, cmd)

    except Exception as e:
        error_msg = str(e)
        # Print collected output to stderr on error
        print(output_buffer.getvalue(), file=sys.stderr)
        logger.error(error_msg)

        if args.json:
            result = {
                'success': False,
                'error': error_msg
            }
            print(json.dumps(result))
        else:
            print(f"Error: {error_msg}", file=sys.stderr)

        return 1

    if args.json:
        result = {
            'success': True,
            'output': output_filename
        }
        print(json.dumps(result))
    else:
        print(f"Screenshot saved to {output_filename}")

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