streamtray.py

from pathlib import Path
import sys, os, threading, time, signal, shutil, uuid, logging, webbrowser, sqlite3
from collections import deque

# ── librerie terze parti ────────────────────────────────────────────────────
from flask import Flask, Response, jsonify, abort, request, redirect, url_for
import cv2
from PIL import Image
import waitress
from pystray import Icon, Menu, MenuItem

TARGET_FPS  = 30
BUF_LEN     = 4          # frame mantenuti in memoria per ogni cam

# ── percorsi e logging ──────────────────────────────────────────────────────
if getattr(sys, "frozen", False):
    RES_DIR  = Path(sys._MEIPASS)
    # Use standard macOS Application Support for persistence
    DATA_DIR = Path.home() / "Library" / "Application Support" / "StreamTray"
else:
    RES_DIR  = Path(__file__).resolve().parent
    DATA_DIR = RES_DIR

DATA_DIR.mkdir(parents=True, exist_ok=True)

ICON_FILE, IMAGE_FILE = RES_DIR / "app.ico", RES_DIR / "app.png"
DB_FILE               = DATA_DIR / "rtsp_streams.db"

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s [%(name)s] %(message)s",
    datefmt="%H:%M:%S")
log = logging.getLogger("streamtray")

# ── DB helpers ───────────────────────────────────────────────────────────────
def get_db():
    DB_FILE.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(DB_FILE); conn.row_factory = sqlite3.Row; return conn

def q(sql, args=()):
    with get_db() as c:
        cur = c.execute(sql, args); c.commit(); return cur.fetchall()

q("""CREATE TABLE IF NOT EXISTS rtsp_streams (
       id INTEGER PRIMARY KEY AUTOINCREMENT,
       camera_id TEXT, rtsp_url TEXT, name TEXT DEFAULT '')""")

# Migration helper: ensure 'name' column exists
try:
    q("SELECT name FROM rtsp_streams LIMIT 1")
except sqlite3.OperationalError:
    log.info("Migrating DB: adding 'name' column")
    q("ALTER TABLE rtsp_streams ADD COLUMN name TEXT DEFAULT ''")

# ── CameraWorker: un thread di cattura per cam ──────────────────────────────
class CameraWorker:
    def __init__(self, url: str):
        self.url  = url
        self.lock = threading.Lock()
        self.views = 0
        self.stop = threading.Event()
        self.th   = None

        # Optimization: Centralized caching
        self.latest_jpeg = None
        self.frame_id    = 0
        self.max_width   = 1280 # Downscale 4K streams for performance

    def _loop(self):
        # UDP-like approach: capture & encode as fast as possible (or capped FPS)
        # Clients just pick the latest. If they are slow, they skip frames.

        # forcing TCP for reliability of transport, but we handle "drops" logically
        url = f"{self.url}?rtsp_transport=tcp&buffer_size=1048576"

        cap = cv2.VideoCapture(url, cv2.CAP_FFMPEG)
        if not cap.isOpened():
            log.error("Impossibile aprire %s", url); return

        fps_delay = 1 / TARGET_FPS

        while not self.stop.is_set():
            t0 = time.time()
            ok, frame = cap.read()
            if not ok:
                log.warning("Frame perso su %s", self.url)
                time.sleep(1) # Prevent busy loop on loss
                # Try reconnect? For now just loop.
                if not cap.isOpened(): break
                continue

            # 1. Resize if huge (Performance)
            h, w = frame.shape[:2]
            if w > self.max_width:
                scale = self.max_width / w
                frame = cv2.resize(frame, (int(w*scale), int(h*scale)), interpolation=cv2.INTER_AREA)

            # 2. Centralized Encoding (Critical Optimization)
            # Encode ONCE here, not in every client thread
            ok, buf = cv2.imencode(".jpg", frame, [int(cv2.IMWRITE_JPEG_QUALITY), 65])
            if ok:
                with self.lock:
                    self.latest_jpeg = buf.tobytes()
                    self.frame_id   += 1

            # Cap FPS to save CPU
            s = fps_delay - (time.time() - t0)
            if s > 0: time.sleep(s)

        cap.release()

    # API --------------------------------------------------------------------
    def start(self):
        if not self.th or not self.th.is_alive():
            self.stop.clear()
            self.th = threading.Thread(target=self._loop, daemon=True); self.th.start()

    def subscribe(self):
        with self.lock:
            self.views += 1
            self.start()

    def unsubscribe(self):
        with self.lock:
            self.views = max(0, self.views - 1)
            if self.views == 0:
                self.stop.set()

    def get_frame(self):
        """Returns (frame_id, jpeg_bytes)"""
        with self.lock:
            return self.frame_id, self.latest_jpeg

# ── registry {camera_id: worker} ────────────────────────────────────────────
_workers, _workers_lock = {}, threading.Lock()

def worker_for(cam_id: str, rtsp_url: str):
    with _workers_lock:
        w = _workers.get(cam_id)
        if not w:
            w = _workers[cam_id] = CameraWorker(rtsp_url)
        # Update URL if changed (basic support)
        if w.url != rtsp_url:
            w.url = rtsp_url
            # If running, logic to restart? For now simple update.
        return w

def reload_workers():
    """Ricarica la cache degli URL dal DB per essere sicuri."""
    load_rtsp_streams_from_db()

# ── Flask app ───────────────────────────────────────────────────────────────
app = Flask(__name__)

# Cache dict {camera_id: {"url": str, "name": str}}
rtsp_streams: dict[str, dict] = {}

def load_rtsp_streams_from_db():
    rtsp_streams.clear()
    for r in q("SELECT camera_id, rtsp_url, name FROM rtsp_streams"):
        rtsp_streams[r["camera_id"]] = {"url": r["rtsp_url"], "name": r["name"]}

load_rtsp_streams_from_db()

@app.route("/")
def index():
    if not rtsp_streams:
         load_rtsp_streams_from_db()
    # Passiamo la lista direttamente al template
    # Structure: [{"id":..., "url":..., "name":...}]
    streams = [{"id": k, "url": v["url"], "name": v["name"]} for k, v in rtsp_streams.items()]
    return flask_render_template("index.html", rtsp_urls=streams)

# Helper per renderizzare template senza import render_template se non necessario o per chiarezza
from flask import render_template as flask_render_template

@app.route("/api/streams", methods=["GET"])
def api_get_streams():
    load_rtsp_streams_from_db()
    return jsonify([{"id": k, "url": v["url"], "name": v["name"]} for k, v in rtsp_streams.items()])

@app.route("/api/streams", methods=["POST"])
def api_add_stream():
    data = request.json
    if not data or "url" not in data:
        return jsonify({"error": "Missing 'url'"}), 400

    url = data["url"].strip()
    name = data.get("name", "").strip() # Optional name

    if not url:
        return jsonify({"error": "Empty URL"}), 400

    cam_id = str(uuid.uuid4())
    q("INSERT INTO rtsp_streams (camera_id, rtsp_url, name) VALUES (?,?,?)", (cam_id, url, name))
    load_rtsp_streams_from_db()
    return jsonify({"id": cam_id, "url": url, "name": name}), 201

@app.route("/api/streams/<cam_id>", methods=["DELETE"])
def api_delete_stream(cam_id):
    q("DELETE FROM rtsp_streams WHERE camera_id=?", (cam_id,))
    load_rtsp_streams_from_db()
    # Fermiamo eventuali worker attivi
    with _workers_lock:
        if cam_id in _workers:
            # Forziamo stop se necessario, o lasciamo che il GC/unsubscribe faccia il suo corso
            # Qui ci limitiamo a rimuoverlo dalla mappa per evitare riutilizzo errato
            _workers.pop(cam_id, None)

    return jsonify({"status": "deleted"})

@app.route("/api/streams/<cam_id>", methods=["PUT"])
def api_update_stream(cam_id):
    data = request.json
    if not data or "url" not in data:
        return jsonify({"error": "Missing 'url'"}), 400

    new_url = data["url"].strip()
    new_name = data.get("name", "").strip()

    q("UPDATE rtsp_streams SET rtsp_url=?, name=? WHERE camera_id=?", (new_url, new_name, cam_id))
    load_rtsp_streams_from_db()

    # Aggiorna worker esistente se c'è
    with _workers_lock:
        if cam_id in _workers:
            _workers[cam_id].url = new_url

    return jsonify({"id": cam_id, "url": new_url, "name": new_name})


@app.route("/api/snapshot/<cam_id>")
def api_snapshot(cam_id):
    if cam_id not in rtsp_streams:
        load_rtsp_streams_from_db()

    if cam_id not in rtsp_streams:
        abort(404, "Camera not found")

    w = worker_for(cam_id, rtsp_streams[cam_id]["url"])

    # "Keep alive" logic: ensure it runs briefly for a snapshot
    w.subscribe()
    try:
        # Wait up to 2s for a frame if just started
        frame_bytes = None
        for _ in range(40):
            _, frame_bytes = w.get_frame()
            if frame_bytes is not None: break
            time.sleep(0.05)

        if frame_bytes is None:
            abort(503, "No frame available")

        return Response(frame_bytes, mimetype="image/jpeg")
    finally:
        w.unsubscribe()

@app.route("/video_feed/<cam_id>")
def video_feed(cam_id):
    if cam_id not in rtsp_streams:
        load_rtsp_streams_from_db()
    if cam_id not in rtsp_streams:
        abort(404, "Camera not found")

    w = worker_for(cam_id, rtsp_streams[cam_id]["url"])
    w.subscribe()

    def stream():
        last_id = -1
        try:
            while True:
                # 1. Get latest frame (O(1) operation)
                fid, jpeg = w.get_frame()

                # 2. Deduplication: only send if new
                if fid > last_id and jpeg:
                    last_id = fid
                    yield (b"--frame\r\nContent-Type: image/jpeg\r\n\r\n" +
                           jpeg + b"\r\n")
                    # No sleep needed here usually as get_frame is non-blocking
                    # IF we are faster than FPS, next loop hits "if fid > last_id" check.

                # 3. Throttle checking slightly to save CPU
                time.sleep(1.0 / (TARGET_FPS + 5))

        except GeneratorExit:
            pass
        finally:
            w.unsubscribe()

    return Response(stream(), mimetype="multipart/x-mixed-replace; boundary=frame")


# ── Tray + server thread ────────────────────────────────────────────────────

def open_dashboard(icon=None, item=None):
    webbrowser.open("http://localhost:5050/")

def run_flask():
    waitress.serve(app, host="0.0.0.0", port=5050, expose_tracebacks=True, threads=6)

def start_background_server():
    t = threading.Thread(target=run_flask, daemon=True)
    t.start()
    return t

# ── bootstrap ───────────────────────────────────────────────────────────────
if __name__ == "__main__":
    # Avvia subito il server in background
    start_background_server()
    log.info("Server started on http://localhost:5050")

    if len(sys.argv) > 1 and sys.argv[1] == "--no-tray":
        # Modalità solo server (utile per debug o docker)
        try:
            while True: time.sleep(1)
        except KeyboardInterrupt:
            pass
    else:
        # Modalità Tray
        Icon("StreamTray", Image.open(IMAGE_FILE), "StreamTray", Menu(
            MenuItem("Open Dashboard", open_dashboard),
            MenuItem("Quit", lambda i, _: (i.stop(), os.kill(os.getpid(), signal.SIGTERM)))
        )).run()