package config

import (
	"fmt"
	"log"

	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

var db *gorm.DB

// InitDB 初始化数据库连接并创建表
func InitDB() error {
	dbPath, err := GetDBPath()
	if err != nil {
		return fmt.Errorf("获取数据库路径失败: %w", err)
	}

	// 打开数据库连接
	db, err = gorm.Open(sqlite.Open(dbPath), &gorm.Config{
		Logger: logger.Default.LogMode(logger.Info),
	})
	if err != nil {
		return fmt.Errorf("连接数据库失败: %w", err)
	}

	// 启用 WAL 模式以提升并发性能
	if err := db.Exec("PRAGMA journal_mode=WAL").Error; err != nil {
		log.Printf("警告: 启用 WAL 模式失败: %v", err)
	}

	// 自动迁移表结构
	if err := db.AutoMigrate(&Provider{}, &Model{}, &UsageStats{}); err != nil {
		return fmt.Errorf("创建表失败: %w", err)
	}

	log.Printf("数据库初始化成功: %s", dbPath)
	return nil
}

// GetDB 获取数据库连接
func GetDB() *gorm.DB {
	return db
}

// CloseDB 关闭数据库连接
func CloseDB() error {
	if db != nil {
		sqlDB, err := db.DB()
		if err != nil {
			return err
		}
		return sqlDB.Close()
	}
	return nil
}