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

import json

import gi

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

Gst.init([])

# Import cv2 and gstgva after Gst.init to avoid a re-initialization conflict.
import cv2
import numpy as np
from gstgva import VideoFrame

MODEL_DIR = "qwen2_vl_2b_ov"
INPUT_VIDEO = "test_video.mp4"
PROMPT = "Is there fire or smoke in this image? Answer with a single word: yes or no."
ALERT_KEYWORDS = ("yes", "fire", "smoke")

# For CPU: change device=GPU to device=CPU.
# Supported devices: CPU and GPU only.
pipeline_str = (
    f"filesrc location={INPUT_VIDEO} ! decodebin3 ! "
    f"videoconvert ! video/x-raw,format=RGB ! "
    f"gvagenai name=genai model-path={MODEL_DIR} device=GPU "
    f'prompt="{PROMPT}" generation-config="max_new_tokens=8" '
    f"frame-rate=2 chunk-size=2 ! queue ! "
    f"videoconvert ! video/x-raw,format=BGR ! "
    f"appsink name=sink emit-signals=true sync=false max-buffers=4 drop=false"
)
pipeline = Gst.parse_launch(pipeline_str)
appsink = pipeline.get_by_name("sink")

state = {"writer": None, "frame": 0, "detected": False}


def on_sample(sink):
    sample = sink.emit("pull-sample")
    if sample is None:
        return Gst.FlowReturn.OK

    buf = sample.get_buffer()
    caps = sample.get_caps().get_structure(0)
    width = caps.get_value("width")
    height = caps.get_value("height")

    # Read the latest VLM answer from the gvagenai JSON metadata (if present).
    for message in VideoFrame(buf).messages():
        try:
            answer = str(json.loads(message).get("result", "")).strip().lower()
        except (ValueError, TypeError):
            continue
        if answer:
            state["detected"] = any(k in answer for k in ALERT_KEYWORDS)

    ok, mapinfo = buf.map(Gst.MapFlags.READ)
    if not ok:
        return Gst.FlowReturn.OK
    frame = np.frombuffer(mapinfo.data, np.uint8).reshape(height, width, 3).copy()
    buf.unmap(mapinfo)

    if state["detected"]:
        text = "FIRE / SMOKE DETECTED"
        color = (0, 0, 255)  # red alert
    else:
        text = "NO FIRE / SMOKE"
        color = (0, 180, 0)  # green

    # Draw the alert banner across the top of the frame (no bounding boxes).
    cv2.rectangle(frame, (0, 0), (width, 60), (0, 0, 0), -1)
    cv2.putText(frame, text, (20, 42),
                cv2.FONT_HERSHEY_SIMPLEX, 1.2, color, 3)

    if state["writer"] is None:
        state["writer"] = cv2.VideoWriter(
            "output_dlstreamer.mp4",
            cv2.VideoWriter_fourcc(*"mp4v"), 30.0, (width, height),
        )
    state["writer"].write(frame)

    state["frame"] += 1
    if state["frame"] % 30 == 0:
        print(f"frame {state['frame']}: {text}", flush=True)
    return Gst.FlowReturn.OK


appsink.connect("new-sample", on_sample)

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)

if state["writer"] is not None:
    state["writer"].release()
print("Saved: output_dlstreamer.mp4")