# SPDX-License-Identifier: MIT
# Copyright (C) Intel Corporation
#
# License Plate Recognition: DLStreamer pipeline sample.
# See README.md for prerequisites and inputs. Run: python3 dlstreamer_sample.py

import os

import gi

gi.require_version("Gst", "1.0")
gi.require_version("GstAnalytics", "1.0")
gi.require_version("DLStreamerWatermarkMeta", "1.0")
from gi.repository import Gst, GLib, GstAnalytics, DLStreamerWatermarkMeta

Gst.init([])

MODELS_DIR = os.path.abspath("./models/yolov8_license_plate_detector")
DETECTOR_XML = (
    f"{MODELS_DIR}/license-plate-reader/models/"
    "yolov8n/yolov8n_retrained.xml"
)
OCR_XML = (
    f"{MODELS_DIR}/license-plate-reader/models/"
    "ch_PP-OCRv4_rec_infer/ch_PP-OCRv4_rec_infer.xml"
)
INPUT_VIDEO = "ParkingVideo.mp4"
DEVICE = "GPU"

pipeline_str = (
    f"filesrc location={INPUT_VIDEO} ! decodebin3 ! "
    f"videoconvert ! queue ! "
    f"gvadetect model={DETECTOR_XML} device={DEVICE} ! queue ! "
    f"videoconvert ! "
    f"gvaclassify model={OCR_XML} device={DEVICE} ! queue ! "
    f"identity name=probe ! "
    f"gvawatermark displ-cfg=show-labels=false ! videoconvert ! video/x-raw,format=I420 ! "
    f"openh264enc ! h264parse ! "
    f"mp4mux ! filesink location=output_dlstreamer.mp4"
)

pipeline = Gst.parse_launch(pipeline_str)


def on_buffer(pad, info):
    buf = info.get_buffer()
    new_buf = buf.copy_deep()
    rmeta = GstAnalytics.buffer_get_analytics_relation_meta(new_buf)
    if not rmeta:
        info.set_buffer(new_buf)
        return Gst.PadProbeReturn.OK
    for od in rmeta.iter_on_type(GstAnalytics.ODMtd):
        _, x, y, w, h, _ = od.get_location()
        # OCR result is attached as ClsMtd with CONTAIN relation
        text = ""
        for cls in od.iter_direct_related(
            GstAnalytics.RelTypes.CONTAIN, GstAnalytics.ClsMtd
        ):
            if cls.get_length() > 0:
                q = cls.get_quark(0)
                text = GLib.quark_to_string(q) if q else ""
            break
        if text:
            DLStreamerWatermarkMeta.text_meta_add(
                new_buf, x=int(x), y=max(0, int(y) - 10),
                text=text, font_scale=0.6, font_type=0,
                r=0, g=255, b=0, thickness=1, draw_bg=True)
            print(f"Plate: {text}  bbox=({x},{y})", flush=True)
    info.set_buffer(new_buf)
    return Gst.PadProbeReturn.OK


probe = pipeline.get_by_name("probe")
probe.get_static_pad("src").add_probe(
    Gst.PadProbeType.BUFFER, on_buffer
)

pipeline.set_state(Gst.State.PLAYING)
bus = pipeline.get_bus()
bus.timed_pop_filtered(
    Gst.CLOCK_TIME_NONE,
    Gst.MessageType.EOS | Gst.MessageType.ERROR,
)
pipeline.set_state(Gst.State.NULL)