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

package database

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

	_ "modernc.org/sqlite"
)

type DB struct {
	conn *sql.Conn
}

type MQTTConfig struct {
	ID       int    `json:"id"`
	Broker   string `json:"broker"`
	ClientID string `json:"client_id"`
	Username string `json:"username,omitempty"`
	Password string `json:"password,omitempty"`
	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("sqlite", 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,
		username TEXT DEFAULT '',
		password TEXT DEFAULT '',
		topic TEXT NOT NULL,
		enabled BOOLEAN DEFAULT 1
	);
	`
	_, err := d.conn.ExecContext(context.Background(), query)
	if err != nil {
		return err
	}

	for _, col := range []string{"username", "password"} {
		alterQuery := fmt.Sprintf("ALTER TABLE mqtt_config ADD COLUMN %s TEXT DEFAULT ''", col)
		d.conn.ExecContext(context.Background(), alterQuery)
	}
	return nil
}

func (d *DB) GetMQTTConfig() (*MQTTConfig, error) {
	query := "SELECT id, broker, client_id, username, password, 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.Username, &cfg.Password, &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, username, password, topic, enabled) VALUES (?, ?, ?, ?, ?, ?)"
		_, err := d.conn.ExecContext(context.Background(), query, cfg.Broker, cfg.ClientID, cfg.Username, cfg.Password, cfg.Topic, cfg.Enabled)
		return err
	}
	query := "UPDATE mqtt_config SET broker = ?, client_id = ?, username = ?, password = ?, topic = ?, enabled = ? WHERE id = ?"
	_, err := d.conn.ExecContext(context.Background(), query, cfg.Broker, cfg.ClientID, cfg.Username, cfg.Password, 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, username, password, 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.Username, &cfg.Password, &cfg.Topic, &cfg.Enabled); err != nil {
			continue
		}
		configs = append(configs, cfg)
	}
	return configs, nil
}

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