#include "hantro_hw.h"
 #include "hantro_v4l2.h"
 
+static int hantro_set_fmt_out(struct hantro_ctx *ctx,
+                             struct v4l2_pix_format_mplane *pix_mp);
+static int hantro_set_fmt_cap(struct hantro_ctx *ctx,
+                             struct v4l2_pix_format_mplane *pix_mp);
+
 static const struct hantro_fmt *
 hantro_get_formats(const struct hantro_ctx *ctx, unsigned int *num_fmts)
 {
        return 0;
 }
 
-static int vidioc_try_fmt(struct file *file, void *priv, struct v4l2_format *f,
-                         bool capture)
+static int hantro_try_fmt(const struct hantro_ctx *ctx,
+                         struct v4l2_pix_format_mplane *pix_mp,
+                         enum v4l2_buf_type type)
 {
-       struct hantro_ctx *ctx = fh_to_ctx(priv);
-       struct v4l2_pix_format_mplane *pix_mp = &f->fmt.pix_mp;
        const struct hantro_fmt *fmt, *vpu_fmt;
+       bool capture = !V4L2_TYPE_IS_OUTPUT(type);
        bool coded;
 
        coded = capture == hantro_is_encoder_ctx(ctx);
        fmt = hantro_find_format(ctx, pix_mp->pixelformat);
        if (!fmt) {
                fmt = hantro_get_default_fmt(ctx, coded);
-               f->fmt.pix_mp.pixelformat = fmt->fourcc;
+               pix_mp->pixelformat = fmt->fourcc;
        }
 
        if (coded) {
 static int vidioc_try_fmt_cap_mplane(struct file *file, void *priv,
                                     struct v4l2_format *f)
 {
-       return vidioc_try_fmt(file, priv, f, true);
+       return hantro_try_fmt(fh_to_ctx(priv), &f->fmt.pix_mp, f->type);
 }
 
 static int vidioc_try_fmt_out_mplane(struct file *file, void *priv,
                                     struct v4l2_format *f)
 {
-       return vidioc_try_fmt(file, priv, f, false);
+       return hantro_try_fmt(fh_to_ctx(priv), &f->fmt.pix_mp, f->type);
 }
 
 static void
        }
 
        hantro_reset_fmt(fmt, vpu_fmt);
-       fmt->num_planes = 1;
        fmt->width = vpu_fmt->frmsize.min_width;
        fmt->height = vpu_fmt->frmsize.min_height;
-       fmt->plane_fmt[0].sizeimage = vpu_fmt->header_size +
-                               fmt->width * fmt->height * vpu_fmt->max_depth;
+       if (hantro_is_encoder_ctx(ctx))
+               hantro_set_fmt_cap(ctx, fmt);
+       else
+               hantro_set_fmt_out(ctx, fmt);
 }
 
 static void
        }
 
        hantro_reset_fmt(raw_fmt, raw_vpu_fmt);
-       v4l2_fill_pixfmt_mp(raw_fmt, raw_vpu_fmt->fourcc,
-                           encoded_fmt->width,
-                           encoded_fmt->height);
+       raw_fmt->width = encoded_fmt->width;
+       raw_fmt->width = encoded_fmt->width;
+       if (hantro_is_encoder_ctx(ctx))
+               hantro_set_fmt_out(ctx, raw_fmt);
+       else
+               hantro_set_fmt_cap(ctx, raw_fmt);
 }
 
 void hantro_reset_fmts(struct hantro_ctx *ctx)
        }
 }
 
-static int
-vidioc_s_fmt_out_mplane(struct file *file, void *priv, struct v4l2_format *f)
+static int hantro_set_fmt_out(struct hantro_ctx *ctx,
+                             struct v4l2_pix_format_mplane *pix_mp)
 {
-       struct v4l2_pix_format_mplane *pix_mp = &f->fmt.pix_mp;
-       struct hantro_ctx *ctx = fh_to_ctx(priv);
-       struct vb2_queue *vq = v4l2_m2m_get_vq(ctx->fh.m2m_ctx, f->type);
+       struct vb2_queue *vq;
        int ret;
 
-       ret = vidioc_try_fmt_out_mplane(file, priv, f);
+       vq = v4l2_m2m_get_vq(ctx->fh.m2m_ctx,
+                            V4L2_BUF_TYPE_VIDEO_OUTPUT_MPLANE);
+       ret = hantro_try_fmt(ctx, pix_mp, V4L2_BUF_TYPE_VIDEO_OUTPUT_MPLANE);
        if (ret)
                return ret;
 
        return 0;
 }
 
-static int vidioc_s_fmt_cap_mplane(struct file *file, void *priv,
-                                  struct v4l2_format *f)
+static int hantro_set_fmt_cap(struct hantro_ctx *ctx,
+                             struct v4l2_pix_format_mplane *pix_mp)
 {
-       struct v4l2_pix_format_mplane *pix_mp = &f->fmt.pix_mp;
-       struct hantro_ctx *ctx = fh_to_ctx(priv);
        struct vb2_queue *vq;
        int ret;
 
        /* Change not allowed if queue is busy. */
-       vq = v4l2_m2m_get_vq(ctx->fh.m2m_ctx, f->type);
+       vq = v4l2_m2m_get_vq(ctx->fh.m2m_ctx,
+                            V4L2_BUF_TYPE_VIDEO_CAPTURE_MPLANE);
        if (vb2_is_busy(vq))
                return -EBUSY;
 
                        return -EBUSY;
        }
 
-       ret = vidioc_try_fmt_cap_mplane(file, priv, f);
+       ret = hantro_try_fmt(ctx, pix_mp, V4L2_BUF_TYPE_VIDEO_CAPTURE_MPLANE);
        if (ret)
                return ret;
 
        return 0;
 }
 
+static int
+vidioc_s_fmt_out_mplane(struct file *file, void *priv, struct v4l2_format *f)
+{
+       return hantro_set_fmt_out(fh_to_ctx(priv), &f->fmt.pix_mp);
+}
+
+static int
+vidioc_s_fmt_cap_mplane(struct file *file, void *priv, struct v4l2_format *f)
+{
+       return hantro_set_fmt_cap(fh_to_ctx(priv), &f->fmt.pix_mp);
+}
+
 const struct v4l2_ioctl_ops hantro_ioctl_ops = {
        .vidioc_querycap = vidioc_querycap,
        .vidioc_enum_framesizes = vidioc_enum_framesizes,