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 Dict, List, Optional, Tuple


@dataclass(frozen=True)
class RouteInfo:
    id: int
    name: 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


@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


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":
        # .../main/xiaozhi-server/plugins_func/nav_copilot/storage.py -> parents[2] == .../main/xiaozhi-server
        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,
                    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,
                    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 ''
                )
                """
            )
            self._conn.commit()

    def create_route(self, name: str) -> int:
        now = int(time.time())
        with self._lock:
            self._conn.execute("UPDATE routes SET active=0")
            cur = self._conn.execute(
                "INSERT INTO routes(name, created_at, active) VALUES(?, ?, 1)",
                (name, 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 add_point(
        self,
        *,
        route_id: int,
        seq: int,
        t: int,
        lat: float,
        lon: float,
        speed_mps: float,
        cum_dist_m: float,
    ) -> None:
        with self._lock:
            self._conn.execute(
                """
                INSERT OR REPLACE INTO route_points(route_id, seq, t, lat, lon, speed_mps, cum_dist_m)
                VALUES(?, ?, ?, ?, ?, ?, ?)
                """,
                (route_id, seq, t, lat, lon, speed_mps, cum_dist_m),
            )
            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 = "",
    ) -> int:
        with self._lock:
            cur = self._conn.execute(
                """
                INSERT INTO hazards(route_id, seq, t, lat, lon, cum_dist_m, type, note)
                VALUES(?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (route_id, seq, t, lat, lon, cum_dist_m, hazard_type, note or ""),
            )
            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.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"]),
                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
                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"]),
            )
            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
                FROM hazards
                WHERE route_id=?
                ORDER BY cum_dist_m 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 ""),
            )
            for row in rows
        ]