ai_fwt_pro/main/xiaozhi-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 RouteInfo:
    id: int
    name: str
    mode: str
    created_at: int
    finished_at: Optional[int]
    points: int
    hazards: int
    active: bool


@dataclass(frozen=True)
class RoutePoint:
    seq: int
    t: int
    lat: float
    lon: float
    speed_mps: float
    cum_dist_m: float
    counter_value: int = 0
    counter_absolute: int = 0


@dataclass(frozen=True)
class Hazard:
    id: int
    route_id: int
    seq: int
    t: int
    lat: float
    lon: float
    cum_dist_m: float
    type: str
    note: str
    counter_value: int = 0
    counter_absolute: int = 0


class NavDB:
    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":
        base = Path(__file__).resolve().parents[2]
        return cls(str(base / "data" / "nav_copilot.sqlite3"))

    def _init_schema(self) -> None:
        with self._lock:
            cur = self._conn.cursor()
            cur.execute(
                """
                CREATE TABLE IF NOT EXISTS routes (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT NOT NULL,
                    mode TEXT NOT NULL DEFAULT 'gps',
                    created_at INTEGER NOT NULL,
                    finished_at INTEGER,
                    active INTEGER NOT NULL DEFAULT 0
                )
                """
            )
            cur.execute(
                """
                CREATE TABLE IF NOT EXISTS route_points (
                    route_id INTEGER NOT NULL,
                    seq INTEGER NOT NULL,
                    t INTEGER NOT NULL,
                    lat REAL NOT NULL,
                    lon REAL NOT NULL,
                    speed_mps REAL NOT NULL,
                    cum_dist_m REAL NOT NULL,
                    counter_value INTEGER NOT NULL DEFAULT 0,
                    counter_absolute INTEGER NOT NULL DEFAULT 0,
                    PRIMARY KEY(route_id, seq)
                )
                """
            )
            cur.execute(
                """
                CREATE TABLE IF NOT EXISTS hazards (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    route_id INTEGER NOT NULL,
                    seq INTEGER NOT NULL,
                    t INTEGER NOT NULL,
                    lat REAL NOT NULL,
                    lon REAL NOT NULL,
                    cum_dist_m REAL NOT NULL,
                    type TEXT NOT NULL,
                    note TEXT NOT NULL DEFAULT '',
                    counter_value INTEGER NOT NULL DEFAULT 0,
                    counter_absolute INTEGER NOT NULL DEFAULT 0
                )
                """
            )

            self._ensure_column(cur, "routes", "mode", "TEXT NOT NULL DEFAULT 'gps'")
            self._ensure_column(
                cur,
                "route_points",
                "counter_value",
                "INTEGER NOT NULL DEFAULT 0",
            )
            self._ensure_column(
                cur,
                "route_points",
                "counter_absolute",
                "INTEGER NOT NULL DEFAULT 0",
            )
            self._ensure_column(
                cur,
                "hazards",
                "counter_value",
                "INTEGER NOT NULL DEFAULT 0",
            )
            self._ensure_column(
                cur,
                "hazards",
                "counter_absolute",
                "INTEGER NOT NULL DEFAULT 0",
            )
            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 create_route(self, name: str, mode: str = "gps") -> int:
        now = int(time.time())
        route_mode = (mode or "gps").strip() or "gps"
        with self._lock:
            self._conn.execute("UPDATE routes SET active=0")
            cur = self._conn.execute(
                "INSERT INTO routes(name, mode, created_at, active) VALUES(?, ?, ?, 1)",
                (name, route_mode, now),
            )
            self._conn.commit()
            return int(cur.lastrowid)

    def finish_route(self, route_id: int) -> None:
        now = int(time.time())
        with self._lock:
            self._conn.execute(
                "UPDATE routes SET finished_at=? WHERE id=?",
                (now, route_id),
            )
            self._conn.commit()

    def set_active_route(self, route_id: int) -> None:
        with self._lock:
            self._conn.execute("UPDATE routes SET active=0")
            self._conn.execute("UPDATE routes SET active=1 WHERE id=?", (route_id,))
            self._conn.commit()

    def get_active_route_id(self) -> Optional[int]:
        with self._lock:
            row = self._conn.execute(
                "SELECT id FROM routes WHERE active=1 ORDER BY id DESC LIMIT 1"
            ).fetchone()
            return int(row["id"]) if row else None

    def get_route_mode(self, route_id: int) -> str:
        with self._lock:
            row = self._conn.execute(
                "SELECT mode FROM routes WHERE id=?",
                (route_id,),
            ).fetchone()
        if not row:
            return "gps"
        return str(row["mode"] or "gps")

    def add_point(
        self,
        *,
        route_id: int,
        seq: int,
        t: int,
        lat: float,
        lon: float,
        speed_mps: float,
        cum_dist_m: float,
        counter_value: int = 0,
        counter_absolute: int = 0,
    ) -> None:
        with self._lock:
            self._conn.execute(
                """
                INSERT OR REPLACE INTO route_points(
                    route_id, seq, t, lat, lon, speed_mps, cum_dist_m, counter_value, counter_absolute
                )
                VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    route_id,
                    seq,
                    t,
                    lat,
                    lon,
                    speed_mps,
                    cum_dist_m,
                    int(counter_value or 0),
                    int(counter_absolute or 0),
                ),
            )
            self._conn.commit()

    def add_hazard(
        self,
        *,
        route_id: int,
        seq: int,
        t: int,
        lat: float,
        lon: float,
        cum_dist_m: float,
        hazard_type: str,
        note: str = "",
        counter_value: int = 0,
        counter_absolute: int = 0,
    ) -> int:
        with self._lock:
            cur = self._conn.execute(
                """
                INSERT INTO hazards(
                    route_id, seq, t, lat, lon, cum_dist_m, type, note, counter_value, counter_absolute
                )
                VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    route_id,
                    seq,
                    t,
                    lat,
                    lon,
                    cum_dist_m,
                    hazard_type,
                    note or "",
                    int(counter_value or 0),
                    int(counter_absolute or 0),
                ),
            )
            self._conn.commit()
            return int(cur.lastrowid)

    def list_routes(self, limit: int = 20) -> List[RouteInfo]:
        with self._lock:
            rows = self._conn.execute(
                """
                SELECT
                    r.id, r.name, r.mode, r.created_at, r.finished_at, r.active,
                    (SELECT COUNT(1) FROM route_points p WHERE p.route_id=r.id) AS points,
                    (SELECT COUNT(1) FROM hazards h WHERE h.route_id=r.id) AS hazards
                FROM routes r
                ORDER BY r.id DESC
                LIMIT ?
                """,
                (int(limit),),
            ).fetchall()
        return [
            RouteInfo(
                id=int(row["id"]),
                name=str(row["name"]),
                mode=str(row["mode"] or "gps"),
                created_at=int(row["created_at"]),
                finished_at=int(row["finished_at"]) if row["finished_at"] else None,
                points=int(row["points"]),
                hazards=int(row["hazards"]),
                active=bool(int(row["active"])),
            )
            for row in rows
        ]

    def load_route_points(self, route_id: int) -> List[RoutePoint]:
        with self._lock:
            rows = self._conn.execute(
                """
                SELECT seq, t, lat, lon, speed_mps, cum_dist_m, counter_value, counter_absolute
                FROM route_points
                WHERE route_id=?
                ORDER BY seq ASC
                """,
                (route_id,),
            ).fetchall()
        return [
            RoutePoint(
                seq=int(row["seq"]),
                t=int(row["t"]),
                lat=float(row["lat"]),
                lon=float(row["lon"]),
                speed_mps=float(row["speed_mps"]),
                cum_dist_m=float(row["cum_dist_m"]),
                counter_value=int(row["counter_value"] or 0),
                counter_absolute=int(row["counter_absolute"] or 0),
            )
            for row in rows
        ]

    def load_hazards(self, route_id: int) -> List[Hazard]:
        with self._lock:
            rows = self._conn.execute(
                """
                SELECT
                    id, route_id, seq, t, lat, lon, cum_dist_m, type, note, counter_value, counter_absolute
                FROM hazards
                WHERE route_id=?
                ORDER BY cum_dist_m ASC, id ASC
                """,
                (route_id,),
            ).fetchall()
        return [
            Hazard(
                id=int(row["id"]),
                route_id=int(row["route_id"]),
                seq=int(row["seq"]),
                t=int(row["t"]),
                lat=float(row["lat"]),
                lon=float(row["lon"]),
                cum_dist_m=float(row["cum_dist_m"]),
                type=str(row["type"]),
                note=str(row["note"] or ""),
                counter_value=int(row["counter_value"] or 0),
                counter_absolute=int(row["counter_absolute"] or 0),
            )
            for row in rows
        ]