ai_fwt_pro/main/xingxing-server/plugins_func/nav_copilot/storage.py

from __future__ import annotations

import os
import sqlite3
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional


@dataclass(frozen=True)
class LineAlert:
    id: int
    device_id: str
    start_lat: float
    start_lon: float
    end_lat: float
    end_lon: float
    coord_type: str
    note: str
    created_at: int


class NavDB:
    _default_instance: Optional["NavDB"] = None
    _default_lock = threading.Lock()

    def __init__(self, db_path: str):
        self.db_path = db_path
        os.makedirs(os.path.dirname(db_path), exist_ok=True)
        self._lock = threading.Lock()
        self._conn = sqlite3.connect(db_path, check_same_thread=False)
        self._conn.row_factory = sqlite3.Row
        self._init_schema()

    @classmethod
    def default(cls) -> "NavDB":
        if cls._default_instance is None:
            with cls._default_lock:
                if cls._default_instance is None:
                    base = Path(__file__).resolve().parents[2]
                    cls._default_instance = cls(str(base / "data" / "nav_copilot.sqlite3"))
        return cls._default_instance

    def _init_schema(self) -> None:
        with self._lock:
            cur = self._conn.cursor()

            # The old route/hazard model is no longer used. Drop it so the
            # SQLite file only contains app-managed line alerts.
            cur.execute("DROP TABLE IF EXISTS route_points")
            cur.execute("DROP TABLE IF EXISTS hazards")
            cur.execute("DROP TABLE IF EXISTS routes")

            cur.execute(
                """
                CREATE TABLE IF NOT EXISTS line_alerts (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    device_id TEXT NOT NULL,
                    start_lat REAL NOT NULL,
                    start_lon REAL NOT NULL,
                    end_lat REAL NOT NULL,
                    end_lon REAL NOT NULL,
                    coord_type TEXT NOT NULL DEFAULT 'BD09',
                    note TEXT NOT NULL,
                    created_at INTEGER NOT NULL
                )
                """
            )
            self._ensure_column(
                cur,
                "line_alerts",
                "coord_type",
                "TEXT NOT NULL DEFAULT 'BD09'",
            )
            cur.execute(
                """
                CREATE INDEX IF NOT EXISTS idx_line_alerts_device_id
                ON line_alerts(device_id)
                """
            )
            self._conn.commit()

    def _ensure_column(
        self,
        cur: sqlite3.Cursor,
        table_name: str,
        column_name: str,
        column_ddl: str,
    ) -> None:
        columns = {
            str(row["name"])
            for row in cur.execute(f"PRAGMA table_info({table_name})").fetchall()
        }
        if column_name not in columns:
            cur.execute(
                f"ALTER TABLE {table_name} ADD COLUMN {column_name} {column_ddl}"
            )

    def count_line_alerts(self, device_id: str) -> int:
        with self._lock:
            row = self._conn.execute(
                "SELECT COUNT(1) AS total FROM line_alerts WHERE device_id=?",
                (device_id,),
            ).fetchone()
        return int(row["total"] or 0) if row else 0

    def get_line_alert(self, alert_id: int, device_id: Optional[str] = None) -> Optional[LineAlert]:
        query = """
            SELECT id, device_id, start_lat, start_lon, end_lat, end_lon, coord_type, note, created_at
            FROM line_alerts
            WHERE id=?
        """
        params: list[object] = [int(alert_id)]
        if device_id is not None:
            query += " AND device_id=?"
            params.append(device_id)

        with self._lock:
            row = self._conn.execute(query, tuple(params)).fetchone()
        return self._row_to_line_alert(row) if row else None

    def list_line_alerts(self, device_id: str) -> List[LineAlert]:
        with self._lock:
            rows = self._conn.execute(
                """
                SELECT id, device_id, start_lat, start_lon, end_lat, end_lon, coord_type, note, created_at
                FROM line_alerts
                WHERE device_id=?
                ORDER BY id DESC
                """,
                (device_id,),
            ).fetchall()
        return [self._row_to_line_alert(row) for row in rows]

    def create_line_alert(
        self,
        *,
        device_id: str,
        start_lat: float,
        start_lon: float,
        end_lat: float,
        end_lon: float,
        coord_type: str,
        note: str,
    ) -> int:
        now = int(time.time())
        with self._lock:
            cur = self._conn.execute(
                """
                INSERT INTO line_alerts(
                    device_id, start_lat, start_lon, end_lat, end_lon, coord_type, note, created_at
                )
                VALUES(?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    device_id,
                    float(start_lat),
                    float(start_lon),
                    float(end_lat),
                    float(end_lon),
                    coord_type,
                    note,
                    now,
                ),
            )
            self._conn.commit()
            return int(cur.lastrowid)

    def delete_line_alert(self, alert_id: int, device_id: Optional[str] = None) -> bool:
        query = "DELETE FROM line_alerts WHERE id=?"
        params: list[object] = [int(alert_id)]
        if device_id is not None:
            query += " AND device_id=?"
            params.append(device_id)

        with self._lock:
            cur = self._conn.execute(query, tuple(params))
            self._conn.commit()
            return int(cur.rowcount or 0) > 0

    @staticmethod
    def _row_to_line_alert(row: sqlite3.Row) -> LineAlert:
        return LineAlert(
            id=int(row["id"]),
            device_id=str(row["device_id"]),
            start_lat=float(row["start_lat"]),
            start_lon=float(row["start_lon"]),
            end_lat=float(row["end_lat"]),
            end_lon=float(row["end_lon"]),
            coord_type=str(row["coord_type"] or "BD09"),
            note=str(row["note"] or ""),
            created_at=int(row["created_at"]),
        )