ai-status-light/internal/database/database.go

package database

import (
	"context"
	"database/sql"
	"fmt"
	"os"
	"path/filepath"

	_ "github.com/mattn/go-sqlite3"
)

type DB struct {
	conn *sql.Conn
}

type MQTTConfig struct {
	ID       int    `json:"id"`
	Broker   string `json:"broker"`
	ClientID string `json:"client_id"`
	Topic    string `json:"topic"`
	Enabled  bool   `json:"enabled"`
}

func New(dbPath string) (*DB, error) {
	dir := filepath.Dir(dbPath)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return nil, fmt.Errorf("创建目录失败: %w", err)
	}

	db, err := sql.Open("sqlite3", dbPath)
	if err != nil {
		return nil, fmt.Errorf("打开数据库失败: %w", err)
	}

	conn, err := db.Conn(context.Background())
	if err != nil {
		return nil, fmt.Errorf("获取连接失败: %w", err)
	}

	d := &DB{conn: conn}
	if err := d.init(); err != nil {
		return nil, err
	}
	return d, nil
}

func (d *DB) init() error {
	query := `
	CREATE TABLE IF NOT EXISTS mqtt_config (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		broker TEXT NOT NULL,
		client_id TEXT NOT NULL,
		topic TEXT NOT NULL,
		enabled BOOLEAN DEFAULT 1
	);
	`
	_, err := d.conn.ExecContext(context.Background(), query)
	return err
}

func (d *DB) GetMQTTConfig() (*MQTTConfig, error) {
	query := "SELECT id, broker, client_id, topic, enabled FROM mqtt_config WHERE enabled = 1 LIMIT 1"
	row := d.conn.QueryRowContext(context.Background(), query)

	var cfg MQTTConfig
	err := row.Scan(&cfg.ID, &cfg.Broker, &cfg.ClientID, &cfg.Topic, &cfg.Enabled)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("查询配置失败: %w", err)
	}
	return &cfg, nil
}

func (d *DB) SaveMQTTConfig(cfg *MQTTConfig) error {
	if cfg.ID == 0 {
		query := "INSERT INTO mqtt_config (broker, client_id, topic, enabled) VALUES (?, ?, ?, ?)"
		_, err := d.conn.ExecContext(context.Background(), query, cfg.Broker, cfg.ClientID, cfg.Topic, cfg.Enabled)
		return err
	}
	query := "UPDATE mqtt_config SET broker = ?, client_id = ?, topic = ?, enabled = ? WHERE id = ?"
	_, err := d.conn.ExecContext(context.Background(), query, cfg.Broker, cfg.ClientID, cfg.Topic, cfg.Enabled, cfg.ID)
	return err
}

func (d *DB) DeleteMQTTConfig(id int) error {
	query := "DELETE FROM mqtt_config WHERE id = ?"
	_, err := d.conn.ExecContext(context.Background(), query, id)
	return err
}

func (d *DB) ListMQTTConfigs() ([]MQTTConfig, error) {
	query := "SELECT id, broker, client_id, topic, enabled FROM mqtt_config ORDER BY id"
	rows, err := d.conn.QueryContext(context.Background(), query)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var configs []MQTTConfig
	for rows.Next() {
		var cfg MQTTConfig
		if err := rows.Scan(&cfg.ID, &cfg.Broker, &cfg.ClientID, &cfg.Topic, &cfg.Enabled); err != nil {
			continue
		}
		configs = append(configs, cfg)
	}
	return configs, nil
}

func (d *DB) Close() error {
	return d.conn.Close()
}