GoTunnel/internal/server/db/sqlite.go

package db

import (
	"database/sql"
	"encoding/json"
	"sync"
	"time"

	_ "modernc.org/sqlite"

	"github.com/gotunnel/pkg/protocol"
)

// SQLiteStore SQLite 存储实现
type SQLiteStore struct {
	db *sql.DB
	mu sync.RWMutex
}

// NewSQLiteStore 创建 SQLite 存储
func NewSQLiteStore(dbPath string) (*SQLiteStore, error) {
	db, err := sql.Open("sqlite", dbPath)
	if err != nil {
		return nil, err
	}

	s := &SQLiteStore{db: db}
	if err := s.init(); err != nil {
		db.Close()
		return nil, err
	}

	return s, nil
}

// init 初始化数据库表
func (s *SQLiteStore) init() error {
	// 创建客户端表
	_, err := s.db.Exec(`
		CREATE TABLE IF NOT EXISTS clients (
			id TEXT PRIMARY KEY,
			nickname TEXT NOT NULL DEFAULT '',
			rules TEXT NOT NULL DEFAULT '[]',
			plugins TEXT NOT NULL DEFAULT '[]'
		)
	`)
	if err != nil {
		return err
	}

	// 迁移：添加 nickname 列
	s.db.Exec(`ALTER TABLE clients ADD COLUMN nickname TEXT NOT NULL DEFAULT ''`)
	// 迁移：添加 plugins 列
	s.db.Exec(`ALTER TABLE clients ADD COLUMN plugins TEXT NOT NULL DEFAULT '[]'`)

	// 创建 JS 插件表
	_, err = s.db.Exec(`
		CREATE TABLE IF NOT EXISTS js_plugins (
			name TEXT PRIMARY KEY,
			source TEXT NOT NULL,
			signature TEXT NOT NULL DEFAULT '',
			description TEXT,
			author TEXT,
			version TEXT DEFAULT '',
			auto_push TEXT NOT NULL DEFAULT '[]',
			config TEXT NOT NULL DEFAULT '{}',
			auto_start INTEGER DEFAULT 1,
			enabled INTEGER DEFAULT 1,
			created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
			updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
		)
	`)
	if err != nil {
		return err
	}

	// 迁移：添加 signature 列
	s.db.Exec(`ALTER TABLE js_plugins ADD COLUMN signature TEXT NOT NULL DEFAULT ''`)
	// 迁移：添加 version 列
	s.db.Exec(`ALTER TABLE js_plugins ADD COLUMN version TEXT DEFAULT ''`)
	// 迁移：添加 updated_at 列
	s.db.Exec(`ALTER TABLE js_plugins ADD COLUMN updated_at DATETIME DEFAULT CURRENT_TIMESTAMP`)

	// 创建流量统计表
	_, err = s.db.Exec(`
		CREATE TABLE IF NOT EXISTS traffic_stats (
			hour_ts INTEGER PRIMARY KEY,
			inbound INTEGER NOT NULL DEFAULT 0,
			outbound INTEGER NOT NULL DEFAULT 0
		)
	`)
	if err != nil {
		return err
	}

	// 创建总流量表
	_, err = s.db.Exec(`
		CREATE TABLE IF NOT EXISTS traffic_total (
			id INTEGER PRIMARY KEY CHECK (id = 1),
			inbound INTEGER NOT NULL DEFAULT 0,
			outbound INTEGER NOT NULL DEFAULT 0
		)
	`)
	if err != nil {
		return err
	}

	// 初始化总流量记录
	s.db.Exec(`INSERT OR IGNORE INTO traffic_total (id, inbound, outbound) VALUES (1, 0, 0)`)

	return nil
}

// Close 关闭数据库连接
func (s *SQLiteStore) Close() error {
	return s.db.Close()
}

// GetAllClients 获取所有客户端
func (s *SQLiteStore) GetAllClients() ([]Client, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	rows, err := s.db.Query(`SELECT id, nickname, rules, plugins FROM clients`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var clients []Client
	for rows.Next() {
		var c Client
		var rulesJSON, pluginsJSON string
		if err := rows.Scan(&c.ID, &c.Nickname, &rulesJSON, &pluginsJSON); err != nil {
			return nil, err
		}
		if err := json.Unmarshal([]byte(rulesJSON), &c.Rules); err != nil {
			c.Rules = []protocol.ProxyRule{}
		}
		if err := json.Unmarshal([]byte(pluginsJSON), &c.Plugins); err != nil {
			c.Plugins = []ClientPlugin{}
		}
		clients = append(clients, c)
	}
	return clients, nil
}

// GetClient 获取单个客户端
func (s *SQLiteStore) GetClient(id string) (*Client, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	var c Client
	var rulesJSON, pluginsJSON string
	err := s.db.QueryRow(`SELECT id, nickname, rules, plugins FROM clients WHERE id = ?`, id).Scan(&c.ID, &c.Nickname, &rulesJSON, &pluginsJSON)
	if err != nil {
		return nil, err
	}
	if err := json.Unmarshal([]byte(rulesJSON), &c.Rules); err != nil {
		c.Rules = []protocol.ProxyRule{}
	}
	if err := json.Unmarshal([]byte(pluginsJSON), &c.Plugins); err != nil {
		c.Plugins = []ClientPlugin{}
	}
	return &c, nil
}

// CreateClient 创建客户端
func (s *SQLiteStore) CreateClient(c *Client) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	rulesJSON, err := json.Marshal(c.Rules)
	if err != nil {
		return err
	}
	pluginsJSON, err := json.Marshal(c.Plugins)
	if err != nil {
		return err
	}
	_, err = s.db.Exec(`INSERT INTO clients (id, nickname, rules, plugins) VALUES (?, ?, ?, ?)`,
		c.ID, c.Nickname, string(rulesJSON), string(pluginsJSON))
	return err
}

// UpdateClient 更新客户端
func (s *SQLiteStore) UpdateClient(c *Client) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	rulesJSON, err := json.Marshal(c.Rules)
	if err != nil {
		return err
	}
	pluginsJSON, err := json.Marshal(c.Plugins)
	if err != nil {
		return err
	}
	_, err = s.db.Exec(`UPDATE clients SET nickname = ?, rules = ?, plugins = ? WHERE id = ?`,
		c.Nickname, string(rulesJSON), string(pluginsJSON), c.ID)
	return err
}

// DeleteClient 删除客户端
func (s *SQLiteStore) DeleteClient(id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	_, err := s.db.Exec(`DELETE FROM clients WHERE id = ?`, id)
	return err
}

// ClientExists 检查客户端是否存在
func (s *SQLiteStore) ClientExists(id string) (bool, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	var count int
	err := s.db.QueryRow(`SELECT COUNT(*) FROM clients WHERE id = ?`, id).Scan(&count)
	return count > 0, err
}

// GetClientRules 获取客户端规则
func (s *SQLiteStore) GetClientRules(id string) ([]protocol.ProxyRule, error) {
	c, err := s.GetClient(id)
	if err != nil {
		return nil, err
	}
	return c.Rules, nil
}

// ========== JS 插件存储方法 ==========

// GetAllJSPlugins 获取所有 JS 插件
func (s *SQLiteStore) GetAllJSPlugins() ([]JSPlugin, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	rows, err := s.db.Query(`
		SELECT name, source, signature, description, author, version, auto_push, config, auto_start, enabled
		FROM js_plugins
	`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var plugins []JSPlugin
	for rows.Next() {
		var p JSPlugin
		var autoPushJSON, configJSON string
		var version sql.NullString
		var autoStart, enabled int
		err := rows.Scan(&p.Name, &p.Source, &p.Signature, &p.Description, &p.Author,
			&version, &autoPushJSON, &configJSON, &autoStart, &enabled)
		if err != nil {
			return nil, err
		}
		p.Version = version.String
		json.Unmarshal([]byte(autoPushJSON), &p.AutoPush)
		json.Unmarshal([]byte(configJSON), &p.Config)
		p.AutoStart = autoStart == 1
		p.Enabled = enabled == 1
		plugins = append(plugins, p)
	}
	return plugins, nil
}

// GetJSPlugin 获取单个 JS 插件
func (s *SQLiteStore) GetJSPlugin(name string) (*JSPlugin, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	var p JSPlugin
	var autoPushJSON, configJSON string
	var version sql.NullString
	var autoStart, enabled int
	err := s.db.QueryRow(`
		SELECT name, source, signature, description, author, version, auto_push, config, auto_start, enabled
		FROM js_plugins WHERE name = ?
	`, name).Scan(&p.Name, &p.Source, &p.Signature, &p.Description, &p.Author,
		&version, &autoPushJSON, &configJSON, &autoStart, &enabled)
	if err != nil {
		return nil, err
	}
	p.Version = version.String
	json.Unmarshal([]byte(autoPushJSON), &p.AutoPush)
	json.Unmarshal([]byte(configJSON), &p.Config)
	p.AutoStart = autoStart == 1
	p.Enabled = enabled == 1
	return &p, nil
}

// SaveJSPlugin 保存 JS 插件
func (s *SQLiteStore) SaveJSPlugin(p *JSPlugin) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	autoPushJSON, _ := json.Marshal(p.AutoPush)
	configJSON, _ := json.Marshal(p.Config)
	autoStart, enabled := 0, 0
	if p.AutoStart {
		autoStart = 1
	}
	if p.Enabled {
		enabled = 1
	}

	_, err := s.db.Exec(`
		INSERT OR REPLACE INTO js_plugins
		(name, source, signature, description, author, version, auto_push, config, auto_start, enabled, updated_at)
		VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
	`, p.Name, p.Source, p.Signature, p.Description, p.Author, p.Version,
		string(autoPushJSON), string(configJSON), autoStart, enabled)
	return err
}

// DeleteJSPlugin 删除 JS 插件
func (s *SQLiteStore) DeleteJSPlugin(name string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	_, err := s.db.Exec(`DELETE FROM js_plugins WHERE name = ?`, name)
	return err
}

// SetJSPluginEnabled 设置 JS 插件启用状态
func (s *SQLiteStore) SetJSPluginEnabled(name string, enabled bool) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	val := 0
	if enabled {
		val = 1
	}
	_, err := s.db.Exec(`UPDATE js_plugins SET enabled = ?, updated_at = CURRENT_TIMESTAMP WHERE name = ?`, val, name)
	return err
}

// UpdateJSPluginConfig 更新 JS 插件配置
func (s *SQLiteStore) UpdateJSPluginConfig(name string, config map[string]string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	configJSON, _ := json.Marshal(config)
	_, err := s.db.Exec(`UPDATE js_plugins SET config = ?, updated_at = CURRENT_TIMESTAMP WHERE name = ?`, string(configJSON), name)
	return err
}

// ========== 流量统计方法 ==========

// getHourTimestamp 获取当前小时的时间戳
func getHourTimestamp() int64 {
	now := time.Now()
	return time.Date(now.Year(), now.Month(), now.Day(), now.Hour(), 0, 0, 0, now.Location()).Unix()
}

// AddTraffic 添加流量记录
func (s *SQLiteStore) AddTraffic(inbound, outbound int64) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	hourTs := getHourTimestamp()

	// 更新小时统计
	_, err := s.db.Exec(`
		INSERT INTO traffic_stats (hour_ts, inbound, outbound) VALUES (?, ?, ?)
		ON CONFLICT(hour_ts) DO UPDATE SET inbound = inbound + ?, outbound = outbound + ?
	`, hourTs, inbound, outbound, inbound, outbound)
	if err != nil {
		return err
	}

	// 更新总流量
	_, err = s.db.Exec(`
		UPDATE traffic_total SET inbound = inbound + ?, outbound = outbound + ? WHERE id = 1
	`, inbound, outbound)
	return err
}

// GetTotalTraffic 获取总流量
func (s *SQLiteStore) GetTotalTraffic() (inbound, outbound int64, err error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	err = s.db.QueryRow(`SELECT inbound, outbound FROM traffic_total WHERE id = 1`).Scan(&inbound, &outbound)
	return
}

// Get24HourTraffic 获取24小时流量
func (s *SQLiteStore) Get24HourTraffic() (inbound, outbound int64, err error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	cutoff := time.Now().Add(-24 * time.Hour).Unix()
	err = s.db.QueryRow(`
		SELECT COALESCE(SUM(inbound), 0), COALESCE(SUM(outbound), 0) 
		FROM traffic_stats WHERE hour_ts >= ?
	`, cutoff).Scan(&inbound, &outbound)
	return
}

// GetHourlyTraffic 获取每小时流量记录
func (s *SQLiteStore) GetHourlyTraffic(hours int) ([]TrafficRecord, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	cutoff := time.Now().Add(-time.Duration(hours) * time.Hour).Unix()
	rows, err := s.db.Query(`
		SELECT hour_ts, inbound, outbound FROM traffic_stats 
		WHERE hour_ts >= ? ORDER BY hour_ts ASC
	`, cutoff)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var records []TrafficRecord
	for rows.Next() {
		var r TrafficRecord
		if err := rows.Scan(&r.Timestamp, &r.Inbound, &r.Outbound); err != nil {
			return nil, err
		}
		records = append(records, r)
	}
	return records, nil
}