FreeDATA/freedata_server/message_system_db_broadcasts.py

from freedata_server.message_system_db_manager import DatabaseManager
from freedata_server.message_system_db_model import BroadcastMessage
from freedata_server.message_system_db_station import DatabaseManagerStations
from sqlalchemy.orm.attributes import flag_modified
from sqlalchemy import or_
from datetime import datetime, timedelta, timezone
from freedata_server import helpers
import base64


class DatabaseManagerBroadcasts(DatabaseManager):
    def __init__(self, ctx):
        super().__init__(ctx)

        self.MAX_ATTEMPTS = 20

        # Fixed fallback expiry (hours) used only if a broadcast message is created
        # without an explicit expires_at. Not user-configurable by design.
        self.DEFAULT_EXPIRY_HOURS = 24

        self.stations_manager = DatabaseManagerStations(self.ctx)

    def process_broadcast_message(
        self,
        id: str,
        origin: str,
        timestamp: float,
        burst_index: int,
        burst_data: str,
        total_bursts: int,
        checksum: str,
        repairing_callsigns: dict = None,
        domain: str = None,
        gridsquare: str = None,
        msg_type: str = None,
        received_at: float = None,
        expires_at: float = None,
        nexttransmission_at: float = None,
        priority: int = 1,
        is_read: bool = True,
        direction: str = None,
        status: str = "queued",
        error_reason: str = None,
    ) -> bool:
        """
        Handles both creation of a new broadcast message and addition of bursts.

        If the message does not exist, it will be created.
        If it exists, the burst will be added.
        When all bursts are present, the final payload will be assembled and CRC checked.
        """
        session = self.get_thread_scoped_session()
        try:
            # Try to find existing message
            msg = session.query(BroadcastMessage).filter_by(id=id).first()

            self.log(f"Broadcast ID: {id}, Burst: {burst_index}, Exists: {'yes' if msg else 'no'}")

            if not msg:
                # Create station and status
                origin_station = self.stations_manager.get_or_create_station(origin, session)
                status_obj = self.get_or_create_status(session, status) if status else None

                if expires_at is None:
                    expires_at = (datetime.now(timezone.utc) + timedelta(hours=self.DEFAULT_EXPIRY_HOURS)).timestamp()
                    self.log(f"No expires_at provided for {id}, defaulting to {self.DEFAULT_EXPIRY_HOURS}h from now")

                print("nexttransmission_at", nexttransmission_at)
                print("received_at", received_at)
                print("timestamp", timestamp)
                print("exires_at", expires_at)

                # New message
                msg = BroadcastMessage(
                    id=id,
                    origin=origin_station.callsign,
                    timestamp=timestamp,
                    repairing_callsigns=repairing_callsigns,
                    domain=domain,
                    gridsquare=gridsquare,
                    priority=priority,
                    is_read=is_read,
                    direction=direction,
                    payload_size=0,
                    payload_data={"bursts": {str(burst_index): burst_data}},
                    msg_type=msg_type,
                    total_bursts=total_bursts,
                    checksum=checksum,
                    received_at=received_at,
                    nexttransmission_at=nexttransmission_at,
                    expires_at=expires_at,
                    status_id=status_obj.id if status_obj else None,
                    error_reason=error_reason,
                )
                session.add(msg)
                self.log(f"Created new broadcast message {id}")
                self.ctx.event_manager.freedata_message_db_change(message_id=id)

            else:
                # Add burst to existing message
                if not msg.payload_data:
                    msg.payload_data = {}

                if "bursts" not in msg.payload_data:
                    msg.payload_data["bursts"] = {}

                msg.payload_data["bursts"][str(burst_index)] = burst_data
                flag_modified(msg, "payload_data")

                self.log(f"Added burst {burst_index} to message {id}")
                self.ctx.event_manager.freedata_message_db_change(message_id=id)
            # Check for final assembly
            received = msg.payload_data["bursts"]
            total = msg.total_bursts

            if total > 0 and len(received) == total and all(str(i) in received for i in range(1, total + 1)):
                ordered = [received[str(i)] for i in range(1, total + 1)]
                final_bytes = b"".join(base64.b64decode(b64part) for b64part in ordered)

                # CRC check
                crc = helpers.get_crc_24(final_bytes).hex()

                if msg.checksum is None:
                    self.log(f"Missing checksum for {id}", isWarning=True)
                elif crc != msg.checksum:
                    self.log(f"Checksum mismatch for {id}: expected {msg.checksum}, got {crc}", isWarning=True)
                    status_obj = self.get_or_create_status(session, "failed_checksum")
                    msg.status_id = status_obj.id

                else:
                    msg.payload_data["final"] = base64.b64encode(final_bytes).decode()
                    msg.payload_size = len(final_bytes)

                    status_obj = self.get_or_create_status(session, "received")
                    msg.status_id = status_obj.id

                    self.log(f"Final payload assembled and verified for {id}")

            session.commit()
            self.ctx.event_manager.freedata_message_db_change(message_id=msg.id)
            return True

        except Exception as e:
            session.rollback()
            self.log(f"Error processing broadcast message {id}: {e}", isWarning=True)
            return False

        finally:
            session.remove()

    def get_all_broadcasts_json(self) -> list:
        """
        Returns all broadcast messages in JSON-serializable dict format.
        """
        session = self.get_thread_scoped_session()
        try:
            messages = session.query(BroadcastMessage).all()
            result = []

            for msg in messages:
                result.append({
                    "id": msg.id,
                    "origin": msg.origin,
                    "timestamp": msg.timestamp if msg.timestamp else None,
                    "repairing_callsigns": msg.repairing_callsigns,
                    "domain": msg.domain,
                    "gridsquare": msg.gridsquare,
                    "frequency": msg.frequency,
                    "priority": msg.priority,
                    "is_read": msg.is_read,
                    "direction": msg.direction,
                    "payload_size": msg.payload_size,
                    "payload_data": msg.payload_data,
                    "msg_type": msg.msg_type,
                    "total_bursts": msg.total_bursts,
                    "checksum": msg.checksum,
                    "received_at": msg.received_at if msg.received_at else None,
                    "expires_at": msg.expires_at if msg.expires_at else None,
                    "nexttransmission_at": msg.nexttransmission_at if msg.nexttransmission_at else None,
                    "status": msg.status.name if msg.status else None,
                    "error_reason": msg.error_reason,
                })

            return result

        except Exception as e:
            self.log(f"Error fetching broadcasts: {e}", isWarning=True)
            return []

        finally:
            session.remove()

    def get_first_queued_message(self):
        session = self.get_thread_scoped_session()
        try:
            now_ts = datetime.now(timezone.utc).timestamp()

            # A broadcast can sit in the transmit queue for a long time if the server
            # was offline (e.g. nexttransmission_at from a backoff schedule set weeks ago).
            # Anything whose expires_at has already passed should never be transmitted -
            # mark it as expired instead of silently sending stale data on the next start.
            expired_status = self.get_or_create_status(session, "expired")
            expired_msgs = (
                session
                .query(BroadcastMessage)
                .filter(
                    BroadcastMessage.direction == "transmit",
                    BroadcastMessage.expires_at.isnot(None),
                    BroadcastMessage.expires_at <= now_ts,
                    BroadcastMessage.status_id != expired_status.id,
                )
                .all()
            )
            for expired_msg in expired_msgs:
                self.log(
                    f"Broadcast {expired_msg.id} expired (expires_at={expired_msg.expires_at} <= now={now_ts}), "
                    f"marking as expired and skipping transmission",
                    isWarning=True,
                )
                expired_msg.status_id = expired_status.id
            if expired_msgs:
                session.commit()
                for expired_msg in expired_msgs:
                    self.ctx.event_manager.freedata_message_db_change(message_id=expired_msg.id)

            message = (
                session
                .query(BroadcastMessage)
                .filter(
                    BroadcastMessage.direction == "transmit",
                    BroadcastMessage.attempts < self.MAX_ATTEMPTS,
                    BroadcastMessage.nexttransmission_at <= now_ts,
                    or_(
                        BroadcastMessage.expires_at.is_(None),
                        BroadcastMessage.expires_at > now_ts,
                    ),
                )
                .order_by(BroadcastMessage.nexttransmission_at.asc())
                .first()
            )

            return message

        except Exception as e:
            session.rollback()
            self.log(f"Error at get_first_queued_message: {e}", isWarning=True)
            return None

        finally:
            session.remove()

    def get_broadcast_domains_json(self) -> dict:
        """
        Returns a JSON-compatible dictionary where each key is a domain (e.g. 'BB1AA-2'),
        and each value is a dict containing message statistics for that domain,
        including an unread_count based on the `is_read` boolean field.
        """
        session = self.get_thread_scoped_session()
        try:
            messages = (
                session
                .query(BroadcastMessage)
                .filter(BroadcastMessage.domain.isnot(None))
                .order_by(BroadcastMessage.timestamp.desc())
                .all()
            )

            result: dict[str, dict] = {}

            for msg in messages:
                domain = msg.domain
                if domain not in result:
                    result[domain] = {
                        "message_count": 1,
                        "unread_count": 0 if msg.is_read else 1,
                        "last_message_id": msg.id,
                        "last_payload": msg.payload_data,
                        "last_message_timestamp": msg.timestamp if msg.timestamp else None,
                        "last_origin": msg.origin,
                    }
                else:
                    result[domain]["message_count"] += 1
                    if not msg.is_read:
                        result[domain]["unread_count"] += 1

            return result

        except Exception as e:
            self.log(f"Error fetching domain summary: {e}", isWarning=True)
            return {}

        finally:
            session.remove()

    def get_broadcasts_per_domain_json(self, domain: str = None) -> dict:
        if domain:
            self.mark_domain_as_read(domain)

        session = self.get_thread_scoped_session()
        try:
            query = session.query(BroadcastMessage).order_by(BroadcastMessage.timestamp.asc())

            if domain:
                query = query.filter(BroadcastMessage.domain == domain)

            messages = query.all()
            result = {}

            for msg in messages:
                d = msg.domain or "unknown"

                if domain and d != domain:
                    continue  # just in case

                if d not in result:
                    result[d] = []

                result[d].append({
                    "id": msg.id,
                    "timestamp": msg.timestamp if msg.timestamp else None,
                    "origin": msg.origin,
                    "gridsquare": msg.gridsquare,
                    "msg_type": msg.msg_type,
                    "payload_size": msg.payload_size,
                    "payload_data": msg.payload_data,
                    "direction": msg.direction,
                    "status": msg.status.name if msg.status else None,
                    "error_reason": msg.error_reason,
                    "received_at": msg.received_at if msg.received_at else None,
                    "nexttransmission_at": msg.nexttransmission_at if msg.nexttransmission_at else None,
                    "expires_at": msg.expires_at if msg.expires_at else None,
                })

            return result

        except Exception as e:
            self.log(f"Error collecting broadcasts for domain '{domain}': {e}", isWarning=True)
            return {}

        finally:
            session.remove()

    def delete_broadcast_message_or_domain(self, id) -> dict:
        session = self.get_thread_scoped_session()
        try:
            msg = session.query(BroadcastMessage).filter_by(id=id).first()
            if msg:
                session.delete(msg)
                session.commit()
                self.log(f"Deleted broadcast message {id}")
                self.ctx.event_manager.freedata_message_db_change(message_id=id)
                return {"status": "success", "deleted": 1, "type": "message", "id": id}

            messages = session.query(BroadcastMessage).filter_by(domain=id).all()
            if messages:
                count = len(messages)
                for m in messages:
                    session.delete(m)
                session.commit()
                self.log(f"Deleted {count} messages from domain '{id}'")
                self.ctx.event_manager.freedata_message_db_change(message_id=id)
                return {"status": "success", "deleted": count, "type": "domain", "domain": id}

            return {"status": "error", "message": f"Neither broadcast message ID '{id}' nor domain found."}

        except Exception as e:
            session.rollback()
            self.log(f"Error deleting broadcast message or domain '{id}': {e}", isWarning=True)
            return {"status": "error", "message": str(e)}

        finally:
            session.remove()

    def check_missing_bursts(self):
        session = self.get_thread_scoped_session()
        try:
            now = datetime.now(timezone.utc)
            one_minute_ago_ts = (now - timedelta(minutes=1)).timestamp()
            now = now.timestamp()
            print(now)
            print(one_minute_ago_ts)
            messages = (
                session
                .query(BroadcastMessage)
                .filter(
                    BroadcastMessage.direction == "receive",
                    BroadcastMessage.received_at < one_minute_ago_ts,
                    BroadcastMessage.total_bursts > 0,
                    BroadcastMessage.expires_at > now,
                )
                .order_by(BroadcastMessage.received_at.asc())
                .all()
            )

            for msg in messages:
                if not msg.payload_data or "bursts" not in msg.payload_data:
                    continue

                # Already complete
                if "final" in msg.payload_data:
                    continue

                # Check next transmission time
                print("nexttransmissionat", msg.nexttransmission_at)
                print("now", now)
                print("expires_at", msg.expires_at)
                print("timestamp", msg.timestamp)
                print("received_at", msg.received_at)
                if now <= msg.nexttransmission_at:
                    dt = datetime.fromtimestamp(msg.nexttransmission_at, timezone.utc)
                    self.log(f"Skip {msg.id}: wait until {int(msg.nexttransmission_at)} ({dt.isoformat()})")

                    continue

                # Max attempts
                if msg.attempts >= self.MAX_ATTEMPTS:
                    self.log(f"Skip {msg.id}: max attempts reached")
                    continue

                bursts = msg.payload_data["bursts"]
                total = msg.total_bursts

                missing = [i for i in range(1, total + 1) if str(i) not in bursts]

                if missing:
                    return {
                        "id": msg.id,
                        "origin": msg.origin,
                        "domain": msg.domain,
                        "missing_bursts": missing,
                        "total_bursts": total,
                        "received_bursts": list(bursts.keys()),
                        "received_at": msg.received_at if msg.received_at else None,
                    }

            return None

        except Exception as e:
            self.log(f"Fehler at check_missing_bursts: {e}", isWarning=True)
            return None

        finally:
            session.remove()

    def get_broadcast_per_id(self, id, get_object=False):
        session = self.get_thread_scoped_session()
        try:
            msg = session.query(BroadcastMessage).filter_by(id=id).first()
            if not msg:
                return None

            if get_object:
                return msg

            return {
                "id": msg.id,
                "origin": msg.origin,
                "timestamp": msg.timestamp if msg.timestamp else None,
                "repairing_callsigns": msg.repairing_callsigns,
                "domain": msg.domain,
                "gridsquare": msg.gridsquare,
                "frequency": msg.frequency,
                "priority": msg.priority,
                "is_read": msg.is_read,
                "direction": msg.direction,
                "payload_size": msg.payload_size,
                "payload_data": msg.payload_data,
                "msg_type": msg.msg_type,
                "total_bursts": msg.total_bursts,
                "checksum": msg.checksum,
                "received_at": msg.received_at if msg.received_at else None,
                "expires_at": msg.expires_at if msg.expires_at else None,
                "nexttransmission_at": msg.nexttransmission_at if msg.nexttransmission_at else None,
                "status": msg.status.name if msg.status else None,
                "error_reason": msg.error_reason,
                "attempts": msg.attempts,
            }

        except Exception as e:
            self.log(f"Error fetching broadcast by id '{id}': {e}", isWarning=True)
            return None

        finally:
            session.remove()

    def increment_attempts(self, message_id: str):
        session = self.get_thread_scoped_session()
        try:
            msg = session.query(BroadcastMessage).filter_by(id=message_id).first()
            if msg:
                msg.attempts = (msg.attempts or 0) + 1
                session.commit()
                self.log(f"Increased attempts for message {message_id} to {msg.attempts}")
            else:
                self.log(f"Message {message_id} not found", isWarning=True)
        except Exception as e:
            session.rollback()
            self.log(f"Error incrementing attempts for {message_id}: {e}", isWarning=True)
        finally:
            session.remove()

    def increment_attempts_and_update_next_transmission(self, message_id: str):
        session = self.get_thread_scoped_session()
        try:
            msg = session.query(BroadcastMessage).filter_by(id=message_id).first()
            if not msg:
                self.log(f"Message {message_id} not found", isWarning=True)
                return

            # Increment attempts
            msg.attempts = (msg.attempts or 0) + 1

            # Define backoff intervals (minutes)
            backoff_minutes = [2, 5, 10, 15, 30, 60, 120, 240, 360, 720, 1440, 2880]
            if msg.attempts <= len(backoff_minutes):
                next_delay = backoff_minutes[msg.attempts - 1]
            else:
                next_delay = 2880  # after max backoff minutes reached in table above

            # Update next transmission
            msg.nexttransmission_at = (datetime.now(timezone.utc) + timedelta(minutes=next_delay)).timestamp()
            print("---------------", msg.nexttransmission_at)
            # Check max attempts
            if msg.attempts >= self.MAX_ATTEMPTS:
                # Mark as failed or set specific status
                status_obj = self.get_or_create_status(session, "max_attempts_reached")
                msg.status_id = status_obj.id
                self.log(f"Max attempts reached for {message_id}, marking as failed.")

            session.commit()
            self.log(
                f"Incremented attempts for {message_id} to {msg.attempts}, next transmission at {msg.nexttransmission_at}"
            )

        except Exception as e:
            session.rollback()
            self.log(f"Error incrementing attempts for {message_id}: {e}", isWarning=True)
        finally:
            session.remove()

    def mark_domain_as_read(self, domain: str) -> int:
        """
        Marks all unread *received* messages in a domain as read.
        Returns the number of updated rows.
        """
        session = self.get_thread_scoped_session()
        try:
            msgs = (
                session
                .query(BroadcastMessage)
                .filter(
                    BroadcastMessage.domain == domain,
                    BroadcastMessage.is_read.is_(False),
                )
                .all()
            )
            count = 0
            for m in msgs:
                m.is_read = True
                count += 1
            if count:
                session.commit()
                self.log(f"Marked {count} messages as read in domain '{domain}'")
                self.ctx.event_manager.freedata_message_db_change(message_id=domain)
            return count
        except Exception as e:
            session.rollback()
            self.log(f"Error marking domain '{domain}' as read: {e}", isWarning=True)
            return 0
        finally:
            session.remove()