import cython
import cython.cimports.libav as lib
from cython.cimports.av.codec.context import CodecContext
from cython.cimports.av.codec.hwaccel import HWAccel
from cython.cimports.av.error import err_check
from cython.cimports.av.frame import Frame
from cython.cimports.av.packet import Packet
from cython.cimports.av.utils import avrational_to_fraction, to_avrational
from cython.cimports.av.video.format import VideoFormat, get_pix_fmt, get_video_format
from cython.cimports.av.video.frame import VideoFrame, alloc_video_frame
from cython.cimports.av.video.reformatter import VideoReformatter
from cython.cimports.libc.stdint import int64_t


@cython.cfunc
@cython.exceptval(check=False)
def _get_hw_format(
    ctx: cython.pointer[lib.AVCodecContext],
    pix_fmts: cython.pointer[cython.const[lib.AVPixelFormat]],
) -> lib.AVPixelFormat:
    # In the case where we requested accelerated decoding, the decoder first calls this function
    # with a list that includes both the hardware format and software formats.
    # First we try to pick the hardware format if it's in the list.
    # However, if the decoder fails to initialize the hardware, it will call this function again,
    # with only software formats in pix_fmts. We return ctx->sw_pix_fmt regardless in this case,
    # because that should be in the candidate list. If not, we are out of ideas anyways.
    private_data: cython.pointer[AVCodecPrivateData] = cython.cast(
        cython.pointer[AVCodecPrivateData], ctx.opaque
    )
    i: cython.int = 0
    while pix_fmts[i] != -1:
        if pix_fmts[i] == private_data.hardware_pix_fmt:
            return pix_fmts[i]
        i += 1
    return (
        ctx.sw_pix_fmt if private_data.allow_software_fallback else lib.AV_PIX_FMT_NONE
    )


@cython.final
@cython.cclass
class VideoCodecContext(CodecContext):
    @cython.cfunc
    def _init(
        self,
        ptr: cython.pointer[lib.AVCodecContext],
        codec: cython.pointer[cython.const[lib.AVCodec]],
        hwaccel: HWAccel | None,
    ):
        CodecContext._init(self, ptr, codec, hwaccel)

        if hwaccel is None:
            return

        if self.is_encoder:
            # Hardware-accelerated encoding. We only attach the device context here;
            # the hardware frames context depends on the final width/height/pixel
            # format (set by the user after add_stream()), so it is built lazily in
            # CodecContext.open() via _setup_encode_hwframes().
            self.hwaccel_ctx = hwaccel.create(self.codec, for_encoding=True)
            self.ptr.hw_device_ctx = lib.av_buffer_ref(self.hwaccel_ctx.ptr)
            return

        try:
            self.hwaccel_ctx = hwaccel.create(self.codec)
            self.ptr.hw_device_ctx = lib.av_buffer_ref(self.hwaccel_ctx.ptr)
            self.ptr.pix_fmt = self.hwaccel_ctx.config.ptr.pix_fmt
            self.ptr.get_format = _get_hw_format
            self._private_data.hardware_pix_fmt = self.hwaccel_ctx.config.ptr.pix_fmt
            self._private_data.allow_software_fallback = (
                self.hwaccel.allow_software_fallback
            )
            self.ptr.opaque = cython.address(self._private_data)
        except NotImplementedError:
            # Some streams may not have a hardware decoder. For example, many action
            # cam videos have a low resolution mjpeg stream, which is usually not
            # compatible with hardware decoders.
            # The user may have passed in a hwaccel because they want to decode the main
            # stream with it, so we shouldn't abort even if we find a stream that can't
            # be HW decoded.
            # If the user wants to make sure hwaccel is actually used, they can check with the
            # is_hwaccel() function on each stream's codec context.
            self.hwaccel_ctx = None

    @cython.cfunc
    def _encode_upload_frame(self, vframe: VideoFrame) -> VideoFrame:
        # Upload a software frame onto the device for hardware-accelerated encoding.
        frames_ctx: cython.pointer[lib.AVHWFramesContext] = cython.cast(
            cython.pointer[lib.AVHWFramesContext], self.ptr.hw_frames_ctx.data
        )

        # If the user already handed us a matching hardware frame, pass it through.
        if vframe.ptr.format == frames_ctx.format:
            return vframe

        # Convert to the frames context's software format and size before uploading,
        # since av_hwframe_transfer_data() does not change pixel format or scale.
        if (
            vframe.ptr.format != frames_ctx.sw_format
            or vframe.ptr.width != frames_ctx.width
            or vframe.ptr.height != frames_ctx.height
        ):
            if not self.reformatter:
                self.reformatter = VideoReformatter()
            vframe = self.reformatter.reformat(
                vframe,
                frames_ctx.width,
                frames_ctx.height,
                get_video_format(
                    frames_ctx.sw_format, frames_ctx.width, frames_ctx.height
                ),
                threads=self.ptr.thread_count,
            )

        hwframe: VideoFrame = alloc_video_frame()
        err_check(lib.av_hwframe_get_buffer(self.ptr.hw_frames_ctx, hwframe.ptr, 0))
        err_check(lib.av_hwframe_transfer_data(hwframe.ptr, vframe.ptr, 0))
        hwframe._copy_internal_attributes(vframe, data_layout=False)
        hwframe._init_user_attributes()

        if hwframe.ptr.pts == lib.AV_NOPTS_VALUE:
            hwframe.ptr.pts = self.ptr.frame_num

        return hwframe

    @cython.cfunc
    def _prepare_frames_for_encode(self, input: Frame | None) -> list:
        if input is None or not input:
            return [None]

        vframe: VideoFrame = input

        # Hardware-accelerated encoding: upload the (software) frame to the device.
        if self.ptr.hw_frames_ctx != cython.NULL:
            return [self._encode_upload_frame(vframe)]

        if (
            vframe.format.pix_fmt != self.pix_fmt
            or vframe.width != self.ptr.width
            or vframe.height != self.ptr.height
        ):
            if not self.reformatter:
                self.reformatter = VideoReformatter()

            vframe = self.reformatter.reformat(
                vframe,
                self.ptr.width,
                self.ptr.height,
                self.format,
                threads=self.ptr.thread_count,
            )

        if vframe.ptr.pts == lib.AV_NOPTS_VALUE:
            vframe.ptr.pts = self.ptr.frame_num

        return [vframe]

    @cython.cfunc
    def _alloc_next_frame(self) -> Frame:
        return alloc_video_frame()

    @cython.cfunc
    def _setup_decoded_frame(self, frame: Frame, packet: Packet):
        CodecContext._setup_decoded_frame(self, frame, packet)
        vframe: VideoFrame = frame
        vframe._init_user_attributes()

    @cython.cfunc
    def _transfer_hwframe(self, frame: Frame):
        if self.hwaccel_ctx is None:
            return frame
        if frame.ptr.format != self.hwaccel_ctx.config.ptr.pix_fmt:
            # If we get a software frame, that means we are in software fallback mode, and don't actually
            # need to transfer.
            return frame

        if self.hwaccel_ctx.is_hw_owned:
            cython.cast(VideoFrame, frame)._device_id = self.hwaccel_ctx.device_id
            return frame

        frame_sw: Frame = self._alloc_next_frame()
        err_check(lib.av_hwframe_transfer_data(frame_sw.ptr, frame.ptr, 0))
        frame_sw._copy_internal_attributes(frame, data_layout=False)
        return frame_sw

    @property
    def format(self):
        return get_video_format(
            cython.cast(lib.AVPixelFormat, self.ptr.pix_fmt),
            self.ptr.width,
            self.ptr.height,
        )

    @format.setter
    def format(self, format: VideoFormat):
        self.ptr.pix_fmt = format.pix_fmt
        self.ptr.width = format.width
        self.ptr.height = format.height

    @property
    def width(self):
        if self.ptr is cython.NULL:
            return 0
        return self.ptr.width

    @width.setter
    def width(self, value: cython.uint):
        self.ptr.width = value

    @property
    def height(self):
        if self.ptr is cython.NULL:
            return 0
        return self.ptr.height

    @height.setter
    def height(self, value: cython.uint):
        self.ptr.height = value

    @property
    def bits_per_coded_sample(self):
        """
        The number of bits per sample in the codedwords. It's mandatory for this to be set for some formats to decode properly.

        Wraps :ffmpeg:`AVCodecContext.bits_per_coded_sample`.

        :type: int
        """
        return self.ptr.bits_per_coded_sample

    @bits_per_coded_sample.setter
    def bits_per_coded_sample(self, value: cython.int):
        if self.is_encoder:
            raise ValueError("Not supported for encoders")

        self.ptr.bits_per_coded_sample = value

    @property
    def pix_fmt(self):
        """
        The pixel format's name.

        :type: str | None
        """
        desc: cython.pointer[cython.const[lib.AVPixFmtDescriptor]] = (
            lib.av_pix_fmt_desc_get(cython.cast(lib.AVPixelFormat, self.ptr.pix_fmt))
        )
        if desc == cython.NULL:
            return None
        return cython.cast(str, desc.name)

    @pix_fmt.setter
    def pix_fmt(self, value):
        self.ptr.pix_fmt = get_pix_fmt(value)

    @property
    def sw_format(self):
        """
        For a hardware context (e.g. ``format.name == "cuda"``), the underlying
        software pixel format (``nv12``, ``p010le``, ...). ``None`` for a regular
        software context.

        :type: VideoFormat | None
        """
        if not self.ptr.hw_frames_ctx:
            if self.ptr.sw_pix_fmt != lib.AV_PIX_FMT_NONE:
                return get_video_format(
                    cython.cast(lib.AVPixelFormat, self.ptr.sw_pix_fmt),
                    self.ptr.width,
                    self.ptr.height,
                )
            return None
        frames_ctx: cython.pointer[lib.AVHWFramesContext] = cython.cast(
            cython.pointer[lib.AVHWFramesContext], self.ptr.hw_frames_ctx.data
        )
        return get_video_format(frames_ctx.sw_format, self.ptr.width, self.ptr.height)

    @sw_format.setter
    def sw_format(self, value):
        self.ptr.sw_pix_fmt = get_pix_fmt(value)

    @property
    def framerate(self):
        """
        The frame rate, in frames per second.

        :type: fractions.Fraction
        """
        return avrational_to_fraction(cython.address(self.ptr.framerate))

    @framerate.setter
    def framerate(self, value):
        to_avrational(value, cython.address(self.ptr.framerate))

    @property
    def rate(self):
        """Another name for :attr:`framerate`."""
        return self.framerate

    @rate.setter
    def rate(self, value):
        self.framerate = value

    @property
    def gop_size(self):
        """
        Sets the number of frames between keyframes. Used only for encoding.

        :type: int
        """
        if self.is_decoder:
            raise RuntimeError("Cannot access 'gop_size' as a decoder")
        return self.ptr.gop_size

    @gop_size.setter
    def gop_size(self, value: cython.int):
        if self.is_decoder:
            raise RuntimeError("Cannot access 'gop_size' as a decoder")
        self.ptr.gop_size = value

    @property
    def sample_aspect_ratio(self):
        return avrational_to_fraction(cython.address(self.ptr.sample_aspect_ratio))

    @sample_aspect_ratio.setter
    def sample_aspect_ratio(self, value):
        to_avrational(value, cython.address(self.ptr.sample_aspect_ratio))

    @property
    def display_aspect_ratio(self):
        dar: lib.AVRational
        lib.av_reduce(
            cython.address(dar.num),
            cython.address(dar.den),
            self.ptr.width * self.ptr.sample_aspect_ratio.num,
            self.ptr.height * self.ptr.sample_aspect_ratio.den,
            1024 * 1024,
        )

        return avrational_to_fraction(cython.address(dar))

    @property
    def has_b_frames(self):
        """
        :type: bool
        """
        return bool(self.ptr.has_b_frames)

    @property
    def reorder_depth(self):
        """Raw ``has_b_frames`` value from FFmpeg (int, not bool).

        After :meth:`flush_buffers`, FFmpeg may reset the internal reorder
        heuristic. Set this to the known reorder depth *after* seeking to
        avoid dropped hierarchical B-frames.
        """
        return self.ptr.has_b_frames

    @reorder_depth.setter
    def reorder_depth(self, value: cython.int):
        self.ptr.has_b_frames = value

    @property
    def coded_width(self):
        """
        :type: int
        """
        return self.ptr.coded_width

    @property
    def coded_height(self):
        """
        :type: int
        """
        return self.ptr.coded_height

    @property
    def color_range(self):
        """
        Describes the signal range of the colorspace.

        Wraps :ffmpeg:`AVFrame.color_range`.

        :type: int
        """
        return self.ptr.color_range

    @color_range.setter
    def color_range(self, value):
        self.ptr.color_range = value

    @property
    def color_primaries(self):
        """
        Describes the RGB/XYZ matrix of the colorspace.

        Wraps :ffmpeg:`AVFrame.color_primaries`.

        :type: int
        """
        return self.ptr.color_primaries

    @color_primaries.setter
    def color_primaries(self, value):
        self.ptr.color_primaries = value

    @property
    def color_trc(self):
        """
        Describes the linearization function (a.k.a. transformation characteristics) of the colorspace.

        Wraps :ffmpeg:`AVFrame.color_trc`.

        :type: int
        """
        return self.ptr.color_trc

    @color_trc.setter
    def color_trc(self, value):
        self.ptr.color_trc = value

    @property
    def colorspace(self):
        """
        Describes the YUV/RGB transformation matrix of the colorspace.

        Wraps :ffmpeg:`AVFrame.colorspace`.

        :type: int
        """
        return self.ptr.colorspace

    @colorspace.setter
    def colorspace(self, value):
        self.ptr.colorspace = value

    @property
    def max_b_frames(self):
        """
        The maximum run of consecutive B frames when encoding a video.

        :type: int
        """
        return self.ptr.max_b_frames

    @max_b_frames.setter
    def max_b_frames(self, value):
        self.ptr.max_b_frames = value

    @property
    def qmin(self):
        """
        The minimum quantiser value of an encoded stream.

        Wraps :ffmpeg:`AVCodecContext.qmin`.

        :type: int
        """
        return self.ptr.qmin

    @qmin.setter
    def qmin(self, value):
        self.ptr.qmin = value

    @property
    def qmax(self):
        """
        The maximum quantiser value of an encoded stream.

        Wraps :ffmpeg:`AVCodecContext.qmax`.

        :type: int
        """
        return self.ptr.qmax

    @qmax.setter
    def qmax(self, value):
        self.ptr.qmax = value