import numpy as np
import ffmpeg
import subprocess
import sys, os, time
from typing import Any
if sys.platform != 'win32':
    import fcntl

class FrameWriter:
    '''Video frame writer to the output video file using ffmpeg.
    
    Args:
        output_file_name (str): output video file name
        size (tuple[int, int]): frame size (width, height)
        fps (int): frames per second
        bitrate (str): bit rate
        qmin (int): min video quantizer scale (VBR)
        qmax (int): max video quantizer scale (VBR)
        pix_fmt (str|tuple[str, str]): output pixel format or tuple of input/output pixel format
        audio (ffmpeg.nodes.FilterableStream): audio stream which is generated by ffmpeg.input
        stdout (bool): if true, writer object outputs the ffmpeg messages to stdout
    
    Attributes:
        video_file_name (str): output video file name
        frame (np.ndarray): Array of frame buffer which shape is (size[1], size[0], 3) and dtype is np.uint8
    '''
    
    video_file_name: str
    frame: np.ndarray
    
    def __init__(self, 
                 output_file_name: str, 
                 size: tuple[int, int] = (1280, 720), 
                 fps: int = 30, 
                 bitrate: str = '10240k', 
                 qmin: int = 2, 
                 qmax: int = 31, 
                 pix_fmt: str|tuple[str, str] = 'yuv420p', 
                 audio: ffmpeg.nodes.FilterableStream|None = None, 
                 stdout: bool = False):
        self.video_file_name = output_file_name
        self.fps = fps
        self.stdout = stdout
        video_width, video_height = size
        output_args: list[Any] = [self.video_file_name]
        if audio:
            output_args.insert(0, audio)
        if isinstance(pix_fmt, tuple) or isinstance(pix_fmt, list):
            input_pix_fmt, output_pix_fmt = pix_fmt
        elif isinstance(pix_fmt, str):
            input_pix_fmt, output_pix_fmt = 'rgb24', pix_fmt
        else:
            raise ValueError('Type of pix_fmt must be str or tuple[str, str]')
        if input_pix_fmt == 'rgb24':
            self.frame = np.zeros((video_height, video_width, 3), dtype=np.uint8)
        elif input_pix_fmt == 'rgb48':
            self.frame = np.zeros((video_height, video_width, 3), dtype=np.uint16)
        else:
            raise ValueError(f'Unsupported input pixel format: {input_pix_fmt}')
        if sys.platform == 'win32':
            self.stdout = False
        args = (
            ffmpeg
            .input('pipe:', format='rawvideo', pix_fmt=input_pix_fmt, r=fps, s=f'{video_width}x{video_height}')
            .filter_('colorspace', 'bt709', iall='bt601-6-625', fast='1')
            .output(*output_args, 
                    sws_flags='spline+accurate_rnd+full_chroma_int', 
                    color_range=1, colorspace=1, color_primaries=1, color_trc=1, 
                    pix_fmt=output_pix_fmt, video_bitrate=bitrate, qmin=qmin, qmax=qmax)
            .overwrite_output()
            .compile()
        )
        self.process = subprocess.Popen(
            args, 
            stdin=subprocess.PIPE, 
            stdout=subprocess.DEVNULL, 
            stderr=subprocess.PIPE if self.stdout else subprocess.DEVNULL
        )
        if self.stdout:
            fl =  fcntl.fcntl(self.process.stderr, fcntl.F_GETFL)
            fcntl.fcntl(self.process.stderr, fcntl.F_SETFL, fl|os.O_NONBLOCK)

    def __enter__(self):
        return self

    def __exit__(self, exception_type, exception_value, traceback):
        self.close()

    def close(self):
        '''Close writer object.
        '''
        try:
            self.process.stdin.close()
            self.process.wait()
            self.process.stdout.close()
            self.process.stderr.close()
            if self.stdout:
                self.__print_ffmpeg_messages()
                print()
        except:
            pass

    def write_frame(self):
        '''Add a frame in the buffer of the writer object to the output video file.
        '''
        self.write(self.frame)

    def write(self, frame):
        '''Add a frame of the external buffer to the output video file.

        Args:
            frame (np.ndarray): an external buffer which shape is (size[1], size[0], 3) and dtype is np.uint8
        '''
        if frame.shape != self.frame.shape:
            raise ValueError(f'Shape of frame parameter should be {self.frame.shape}, but {frame.shape}')
        self.process.stdin.write(frame.tobytes())
        self.__print_ffmpeg_messages()
    
    def __print_ffmpeg_messages(self):
        if self.stdout and self.process.stderr.readable():
            data = self.process.stderr.read()
            if data:
                print(data.decode('utf-8'), end='')
            sys.stdout.flush()
