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

package database

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

	_ "modernc.org/sqlite"

	"ai-status-light/internal/logger"
)

type DB struct {
	db   *sql.DB
	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"`
}

type BLEConfig struct {
	ID             int    `json:"id"`
	DeviceName     string `json:"device_name"`
	ServiceUUID    string `json:"service_uuid"`
	ModeCharUUID   string `json:"mode_char_uuid"`
	ConfigCharUUID string `json:"config_char_uuid"`
	Enabled        bool   `json:"enabled"`
}

type DeviceConfig struct {
	ID          int    `json:"id"`
	DeviceName  string `json:"device_name"`
	ConfigTopic string `json:"config_topic"`
	WifiSSID    string `json:"wifi_ssid,omitempty"`
	WifiPass    string `json:"wifi_pass,omitempty"`
	MqttBroker  string `json:"mqtt_broker,omitempty"`
	MqttPort    int    `json:"mqtt_port,omitempty"`
	MqttUser    string `json:"mqtt_user,omitempty"`
	MqttPass    string `json:"mqtt_pass,omitempty"`
	MqttClient  string `json:"mqtt_client,omitempty"`
	MqttTopic   string `json:"mqtt_topic,omitempty"`
	MqttStatus  string `json:"mqtt_status,omitempty"`
	PinRed      int    `json:"pin_red,omitempty"`
	PinGreen    int    `json:"pin_green,omitempty"`
	PinYellow   int    `json:"pin_yellow,omitempty"`
	Enabled     bool   `json:"enabled"`
}

type StatusRecord struct {
	ID        int    `json:"id"`
	Code      string `json:"code"`
	Timestamp string `json:"timestamp"`
}

type WorkDuration struct {
	DurationMinutes int `json:"duration_minutes"`
}

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

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

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

	d := &DB{db: db, conn: conn}
	if err := d.init(); err != nil {
		logger.Error("初始化数据库失败: %v", err)
		return nil, err
	}
	logger.Debug("数据库已打开: %s", dbPath)
	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)
	}

	bleQuery := `
	CREATE TABLE IF NOT EXISTS ble_config (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		device_name TEXT NOT NULL,
		service_uuid TEXT NOT NULL,
		mode_char_uuid TEXT NOT NULL DEFAULT '',
		config_char_uuid TEXT NOT NULL DEFAULT '',
		enabled BOOLEAN DEFAULT 1
	);
	`
	_, err = d.conn.ExecContext(context.Background(), bleQuery)
	if err != nil {
		return err
	}

	for _, col := range []string{"service_uuid", "mode_char_uuid", "config_char_uuid"} {
		alterQuery := fmt.Sprintf("ALTER TABLE ble_config ADD COLUMN %s TEXT NOT NULL DEFAULT ''", col)
		d.conn.ExecContext(context.Background(), alterQuery)
	}

	deviceQuery := `
	CREATE TABLE IF NOT EXISTS device_config (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		device_name TEXT NOT NULL,
		config_topic TEXT NOT NULL DEFAULT 'agent/status/config',
		wifi_ssid TEXT DEFAULT '',
		wifi_pass TEXT DEFAULT '',
		mqtt_broker TEXT DEFAULT '',
		mqtt_port INTEGER DEFAULT 1883,
		mqtt_user TEXT DEFAULT '',
		mqtt_pass TEXT DEFAULT '',
		mqtt_client TEXT DEFAULT '',
		mqtt_topic TEXT DEFAULT '',
		mqtt_status TEXT DEFAULT '',
		pin_red INTEGER DEFAULT 4,
		pin_green INTEGER DEFAULT 3,
		pin_yellow INTEGER DEFAULT 2,
		enabled BOOLEAN DEFAULT 1
	);
	`
	_, err = d.conn.ExecContext(context.Background(), deviceQuery)
	if err != nil {
		return err
	}

	historyQuery := `
	CREATE TABLE IF NOT EXISTS status_history (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		code TEXT NOT NULL,
		timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
	);
	`
	_, err = d.conn.ExecContext(context.Background(), historyQuery)
	if err != nil {
		return err
	}

	_, err = d.conn.ExecContext(context.Background(), "CREATE INDEX IF NOT EXISTS idx_status_history_timestamp ON status_history(timestamp)")
	if err != nil {
		return err
	}

	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 {
		logger.Debug("未找到启用的 MQTT 配置")
		return nil, nil
	}
	if err != nil {
		logger.Error("查询 MQTT 配置失败: %v", err)
		return nil, fmt.Errorf("查询配置失败: %w", err)
	}
	logger.Debug("获取到 MQTT 配置: id=%d, broker=%s", cfg.ID, cfg.Broker)
	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)
		if err != nil {
			logger.Error("插入 MQTT 配置失败: %v", err)
		}
		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)
	if err != nil {
		logger.Error("更新 MQTT 配置失败: id=%d, %v", cfg.ID, err)
	}
	return err
}

func (d *DB) DeleteMQTTConfig(id int) error {
	query := "DELETE FROM mqtt_config WHERE id = ?"
	_, err := d.conn.ExecContext(context.Background(), query, id)
	if err != nil {
		logger.Error("删除 MQTT 配置失败: id=%d, %v", id, err)
	} else {
		logger.Debug("MQTT 配置已删除: id=%d", 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 {
		logger.Error("查询 MQTT 配置列表失败: %v", err)
		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 {
			logger.Warn("扫描 MQTT 配置行失败: %v", err)
			continue
		}
		configs = append(configs, cfg)
	}
	logger.Debug("查询到 %d 条 MQTT 配置", len(configs))
	return configs, nil
}

func (d *DB) Close() error {
	logger.Debug("数据库连接已关闭")
	if err := d.conn.Close(); err != nil {
		return err
	}
	return d.db.Close()
}

func (d *DB) GetBLEConfig() (*BLEConfig, error) {
	query := "SELECT id, device_name, service_uuid, mode_char_uuid, config_char_uuid, enabled FROM ble_config WHERE enabled = 1 LIMIT 1"
	row := d.conn.QueryRowContext(context.Background(), query)

	var cfg BLEConfig
	err := row.Scan(&cfg.ID, &cfg.DeviceName, &cfg.ServiceUUID, &cfg.ModeCharUUID, &cfg.ConfigCharUUID, &cfg.Enabled)
	if err == sql.ErrNoRows {
		logger.Debug("未找到启用的 BLE 配置")
		return nil, nil
	}
	if err != nil {
		logger.Error("查询 BLE 配置失败: %v", err)
		return nil, fmt.Errorf("查询配置失败: %w", err)
	}
	logger.Debug("获取到 BLE 配置: id=%d, device=%s", cfg.ID, cfg.DeviceName)
	return &cfg, nil
}

func (d *DB) SaveBLEConfig(cfg *BLEConfig) error {
	if cfg.ID == 0 {
		query := "INSERT INTO ble_config (device_name, service_uuid, mode_char_uuid, config_char_uuid, enabled) VALUES (?, ?, ?, ?, ?)"
		_, err := d.conn.ExecContext(context.Background(), query, cfg.DeviceName, cfg.ServiceUUID, cfg.ModeCharUUID, cfg.ConfigCharUUID, cfg.Enabled)
		if err != nil {
			logger.Error("插入 BLE 配置失败: %v", err)
		}
		return err
	}
	query := "UPDATE ble_config SET device_name = ?, service_uuid = ?, mode_char_uuid = ?, config_char_uuid = ?, enabled = ? WHERE id = ?"
	_, err := d.conn.ExecContext(context.Background(), query, cfg.DeviceName, cfg.ServiceUUID, cfg.ModeCharUUID, cfg.ConfigCharUUID, cfg.Enabled, cfg.ID)
	if err != nil {
		logger.Error("更新 BLE 配置失败: id=%d, %v", cfg.ID, err)
	}
	return err
}

func (d *DB) DeleteBLEConfig(id int) error {
	query := "DELETE FROM ble_config WHERE id = ?"
	_, err := d.conn.ExecContext(context.Background(), query, id)
	if err != nil {
		logger.Error("删除 BLE 配置失败: id=%d, %v", id, err)
	} else {
		logger.Debug("BLE 配置已删除: id=%d", id)
	}
	return err
}

func (d *DB) ListBLEConfigs() ([]BLEConfig, error) {
	query := "SELECT id, device_name, service_uuid, mode_char_uuid, config_char_uuid, enabled FROM ble_config ORDER BY id"
	rows, err := d.conn.QueryContext(context.Background(), query)
	if err != nil {
		logger.Error("查询 BLE 配置列表失败: %v", err)
		return nil, err
	}
	defer rows.Close()

	var configs []BLEConfig
	for rows.Next() {
		var cfg BLEConfig
		if err := rows.Scan(&cfg.ID, &cfg.DeviceName, &cfg.ServiceUUID, &cfg.ModeCharUUID, &cfg.ConfigCharUUID, &cfg.Enabled); err != nil {
			logger.Warn("扫描 BLE 配置行失败: %v", err)
			continue
		}
		configs = append(configs, cfg)
	}
	logger.Debug("查询到 %d 条 BLE 配置", len(configs))
	return configs, nil
}

func (d *DB) GetDeviceConfig() (*DeviceConfig, error) {
	query := "SELECT id, device_name, config_topic, wifi_ssid, wifi_pass, mqtt_broker, mqtt_port, mqtt_user, mqtt_pass, mqtt_client, mqtt_topic, mqtt_status, pin_red, pin_green, pin_yellow, enabled FROM device_config WHERE enabled = 1 LIMIT 1"
	row := d.conn.QueryRowContext(context.Background(), query)

	var cfg DeviceConfig
	err := row.Scan(&cfg.ID, &cfg.DeviceName, &cfg.ConfigTopic, &cfg.WifiSSID, &cfg.WifiPass, &cfg.MqttBroker, &cfg.MqttPort, &cfg.MqttUser, &cfg.MqttPass, &cfg.MqttClient, &cfg.MqttTopic, &cfg.MqttStatus, &cfg.PinRed, &cfg.PinGreen, &cfg.PinYellow, &cfg.Enabled)
	if err == sql.ErrNoRows {
		logger.Debug("未找到启用的设备配置")
		return nil, nil
	}
	if err != nil {
		logger.Error("查询设备配置失败: %v", err)
		return nil, fmt.Errorf("查询配置失败: %w", err)
	}
	logger.Debug("获取到设备配置: id=%d, device=%s", cfg.ID, cfg.DeviceName)
	return &cfg, nil
}

func (d *DB) SaveDeviceConfig(cfg *DeviceConfig) error {
	if cfg.ID == 0 {
		query := "INSERT INTO device_config (device_name, config_topic, wifi_ssid, wifi_pass, mqtt_broker, mqtt_port, mqtt_user, mqtt_pass, mqtt_client, mqtt_topic, mqtt_status, pin_red, pin_green, pin_yellow, enabled) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"
		_, err := d.conn.ExecContext(context.Background(), query, cfg.DeviceName, cfg.ConfigTopic, cfg.WifiSSID, cfg.WifiPass, cfg.MqttBroker, cfg.MqttPort, cfg.MqttUser, cfg.MqttPass, cfg.MqttClient, cfg.MqttTopic, cfg.MqttStatus, cfg.PinRed, cfg.PinGreen, cfg.PinYellow, cfg.Enabled)
		if err != nil {
			logger.Error("插入设备配置失败: %v", err)
		}
		return err
	}
	query := "UPDATE device_config SET device_name = ?, config_topic = ?, wifi_ssid = ?, wifi_pass = ?, mqtt_broker = ?, mqtt_port = ?, mqtt_user = ?, mqtt_pass = ?, mqtt_client = ?, mqtt_topic = ?, mqtt_status = ?, pin_red = ?, pin_green = ?, pin_yellow = ?, enabled = ? WHERE id = ?"
	_, err := d.conn.ExecContext(context.Background(), query, cfg.DeviceName, cfg.ConfigTopic, cfg.WifiSSID, cfg.WifiPass, cfg.MqttBroker, cfg.MqttPort, cfg.MqttUser, cfg.MqttPass, cfg.MqttClient, cfg.MqttTopic, cfg.MqttStatus, cfg.PinRed, cfg.PinGreen, cfg.PinYellow, cfg.Enabled, cfg.ID)
	if err != nil {
		logger.Error("更新设备配置失败: id=%d, %v", cfg.ID, err)
	}
	return err
}

func (d *DB) DeleteDeviceConfig(id int) error {
	query := "DELETE FROM device_config WHERE id = ?"
	_, err := d.conn.ExecContext(context.Background(), query, id)
	if err != nil {
		logger.Error("删除设备配置失败: id=%d, %v", id, err)
	} else {
		logger.Debug("设备配置已删除: id=%d", id)
	}
	return err
}

func (d *DB) ListDeviceConfigs() ([]DeviceConfig, error) {
	query := "SELECT id, device_name, config_topic, wifi_ssid, wifi_pass, mqtt_broker, mqtt_port, mqtt_user, mqtt_pass, mqtt_client, mqtt_topic, mqtt_status, pin_red, pin_green, pin_yellow, enabled FROM device_config ORDER BY id"
	rows, err := d.conn.QueryContext(context.Background(), query)
	if err != nil {
		logger.Error("查询设备配置列表失败: %v", err)
		return nil, err
	}
	defer rows.Close()

	var configs []DeviceConfig
	for rows.Next() {
		var cfg DeviceConfig
		if err := rows.Scan(&cfg.ID, &cfg.DeviceName, &cfg.ConfigTopic, &cfg.WifiSSID, &cfg.WifiPass, &cfg.MqttBroker, &cfg.MqttPort, &cfg.MqttUser, &cfg.MqttPass, &cfg.MqttClient, &cfg.MqttTopic, &cfg.MqttStatus, &cfg.PinRed, &cfg.PinGreen, &cfg.PinYellow, &cfg.Enabled); err != nil {
			logger.Warn("扫描设备配置行失败: %v", err)
			continue
		}
		configs = append(configs, cfg)
	}
	logger.Debug("查询到 %d 条设备配置", len(configs))
	return configs, nil
}

func (d *DB) SaveStatusRecord(code string) error {
	query := "INSERT INTO status_history (code, timestamp) VALUES (?, datetime('now'))"
	_, err := d.conn.ExecContext(context.Background(), query, code)
	if err != nil {
		logger.Error("保存状态记录失败: %v", err)
	}
	return err
}

func (d *DB) GetStatusHistory(limit int) ([]StatusRecord, error) {
	if limit <= 0 {
		limit = 100
	}
	query := "SELECT id, code, timestamp FROM status_history ORDER BY timestamp DESC LIMIT ?"
	rows, err := d.conn.QueryContext(context.Background(), query, limit)
	if err != nil {
		logger.Error("查询状态历史失败: %v", err)
		return nil, err
	}
	defer rows.Close()

	var records []StatusRecord
	for rows.Next() {
		var r StatusRecord
		if err := rows.Scan(&r.ID, &r.Code, &r.Timestamp); err != nil {
			logger.Warn("扫描状态记录行失败: %v", err)
			continue
		}
		records = append(records, r)
	}
	logger.Debug("查询到 %d 条状态记录", len(records))
	return records, nil
}

func (d *DB) CleanOldStatusRecords(hours int) error {
	if hours <= 0 {
		hours = 2
	}
	query := "DELETE FROM status_history WHERE timestamp < datetime('now', ?)"
	param := fmt.Sprintf("-%d hours", hours)
	result, err := d.conn.ExecContext(context.Background(), query, param)
	if err != nil {
		logger.Error("清理旧状态记录失败: %v", err)
		return err
	}
	rowsAffected, _ := result.RowsAffected()
	if rowsAffected > 0 {
		logger.Info("已清理 %d 条超过 %d 小时的状态记录", rowsAffected, hours)
	}
	return nil
}

func (d *DB) GetTodayWorkDuration() (*WorkDuration, error) {
	query := "SELECT code, timestamp FROM status_history WHERE timestamp >= date('now', 'localtime') ORDER BY timestamp ASC"
	rows, err := d.conn.QueryContext(context.Background(), query)
	if err != nil {
		logger.Error("查询今日工作时长失败: %v", err)
		return nil, err
	}
	defer rows.Close()

	type record struct {
		Code string
		Time time.Time
	}
	var records []record
	for rows.Next() {
		var code, ts string
		if err := rows.Scan(&code, &ts); err != nil {
			continue
		}
		t, err := time.ParseInLocation("2006-01-02 15:04:05", ts, time.Local)
		if err != nil {
			logger.Warn("解析时间戳失败: %s, %v", ts, err)
			continue
		}
		records = append(records, record{Code: code, Time: t})
	}

	active := map[string]bool{
		"busy": true, "reasoning": true, "using_tool": true,
		"running": true, "pending": true, "retry": true,
	}

	var totalMs int64
	now := time.Now()
	for i, r := range records {
		if !active[r.Code] {
			continue
		}
		if i+1 < len(records) {
			totalMs += records[i+1].Time.Sub(r.Time).Milliseconds()
		} else if active[r.Code] {
			totalMs += now.Sub(r.Time).Milliseconds()
		}
	}

	logger.Debug("今日工作时长: %d 分钟", int(totalMs/60000))
	return &WorkDuration{DurationMinutes: int(totalMs / 60000)}, nil
}