TsDB2/DB/core/db.go

package core

import (
	"TsDB2/DB/core/index"
	"TsDB2/DB/core/page"
	"TsDB2/DB/core/types"
	"TsDB2/DB/core/util"
	"encoding/binary"
	"fmt"
	"sync"
	"time"
	"unsafe"
)

type TableMeta struct {
	Name       string
	FieldCount uint32
	TotalRows  uint64
	FirstHeap  uint32
	LastHeap   uint32
	FirstIndex uint32
	LastIndex  uint32
	FirstData  uint32
	LastData   uint32
	RowSize    uint16
}

type DB struct {
	header *HeaderPage
	tables *TablePage

	file *File

	// 防抖自动提交
	autoCommitDelay time.Duration
	autoCommitTimer *time.Timer
	autoCommitMu    sync.Mutex

	// 当前活跃事务（nil 表示无事务）
	currentTxn *Transaction
}

func Open(path string) (*DB, error) {
	ff, err := OpenFile(path)
	if err != nil {
		return nil, err
	}

	// 启用页面缓冲池
	ff.InitBufferPool(DefaultBufferPoolCapacity)

	stat, err := ff.f.Stat()
	if err != nil {
		return nil, err
	}

	var header *HeaderPage
	if stat.Size() == 0 {
		header, err = newHeader("tsdb")
		if err != nil {
			return nil, err
		}
		if err := ff.WriteHeader(header); err != nil {
			return nil, err
		}
	} else {
		header, err = ff.ReadHeader()
		if err != nil {
			return nil, err
		}
	}
	// ==========================
	// 初始化表页
	// ==========================
	tablePage := &TablePage{
		handle: page.HandleDb{
			NewPage: func() uint32 {
				pid := header.newPageWithReader(ff.readPageDisk) // 必须直接读盘, 绕过 BufferPool
				_ = ff.WriteHeader(header)
				return pid
			},
			FreePage: func(pageID uint32) {
				header.recyclePage(pageID, ff.WritePageDirect) // 必须直接写盘! 走 BufferPool 会被 ClearCache 清掉
				_ = ff.WriteHeader(header)
			},
			ReadPage: ff.ReadPage,
			SavePage: ff.WritePage,
		},
	}

	// ==========================
	// 初始化索引元数据页
	// ==========================

	saveHeader := false
	//如果索引页和表页都是空的话则直接创建
	if header.TablePage == 0 {
		header.TablePage = header.newPage()
		tablePage.InitHeader(header.TablePage)
		saveHeader = true
	}

	tablePage.Init(header.TablePage)
	if saveHeader {
		err = ff.WriteHeader(header)
		if err != nil {
			return nil, err
		}
	}

	// ==========================
	// WAL 恢复: 检查是否有未完成的 WAL 事务
	// ==========================
	if err := RecoverWAL(ff, header); err != nil {
		return nil, err
	}

	db := &DB{
		file:            ff,
		header:          header,
		tables:          tablePage,
		autoCommitDelay: DefaultAutoCommitDelay,
	}
	db.startAutoCommit()

	// 设置事务回调（在 BufferPool.WritePage 时缓存原始数据）
	ff.SetOnPageWriteCallback(func(pageID uint32) {
		// 如果有活跃事务，缓存原始页数据
		if db.currentTxn != nil {
			_ = db.currentTxn.cacheOriginalPage(pageID)
		}
	})

	// 脏页阈值: 超过时自动触发 Commit 刷盘, 防止内存暴涨
	ff.SetDirtyThreshold(DefaultDirtyThreshold, func() {
		_ = db.Commit()
	})
	return db, nil
}

func (db *DB) updateTime() {
	db.header.ModifyTime = time.Now().UnixNano()
	_ = db.file.WriteHeader(db.header)
}

func (db *DB) FindTable(name string) (*Table, bool) {
	tbl, ok := db.tables.Find(name)
	if ok && tbl.onWrite == nil {
		tbl.SetOnWrite(db.NotifyWrite)
	}
	return tbl, ok
}

func (db *DB) CreateTable(name string, fields []types.FieldInfo) (*Table, error) {
	tbl, err := db.tables.Create(name, fields)
	if err != nil {
		return nil, err
	}
	tbl.SetOnWrite(db.NotifyWrite)
	db.header.TableCount++
	_ = db.file.WriteHeader(db.header)
	db.updateTime()
	return tbl, nil
}

func (db *DB) DropTable(name string) error {
	// 先清除 BufferPool 缓存, 确保回收时读取磁盘最新数据
	db.file.ClearCache()
	if err := db.tables.Drop(name); err != nil {
		return err
	}
	db.header.TableCount--
	_ = db.file.WriteHeader(db.header)
	db.updateTime()
	// 先刷盘再清缓存, 防止 Drop 中标记的 Flag(0x80) 丢失
	db.file.FlushBufferPool()
	db.file.ClearCache()
	return nil
}

func (db *DB) GetAllTables() []TableMeta {
	res := make([]TableMeta, 0, len(db.tables.tables))
	for _, t := range db.tables.tables {
		if util.IsHighBitSet(t.info.Flag) {
			continue
		}
		res = append(res, TableMeta{
			Name:       util.BytesToString(t.info.Name[:]),
			FieldCount: uint32(t.info.FieldCount),
			TotalRows:  t.info.TotalRows,
			FirstHeap:  t.info.FirstHeap,
			LastHeap:   t.info.LastHeap,
			FirstData:  t.info.FirstData,
			LastData:   t.info.LastData,
			RowSize:    t.info.RowSize,
		})
	}
	return res
}

func (db *DB) Close() error {
	// 停止防抖定时器
	db.stopAutoCommit()
	// 先提交剩余脏页 (走 WAL)
	_ = db.Commit()
	// BufferPool 安全网 (理论上已无脏页)
	db.file.FlushBufferPool()
	return db.file.Close()
}

// Commit 显式提交 — 通过 WAL 保证原子性
// 同时重置防抖定时器
func (db *DB) Commit() error {
	// 取消待执行的自动提交
	db.autoCommitMu.Lock()
	if db.autoCommitTimer != nil {
		db.autoCommitTimer.Stop()
	}
	db.autoCommitMu.Unlock()

	// 1. 收集所有 BufferPool 脏页 (TablePage + IndexPage + DataPage 等)
	dirty := db.file.CollectDirty()
	if len(dirty) == 0 {
		return nil // 无脏页, 跳过
	}

	// 2. 创建事务
	wal := db.BeginTxn()

	// 3. 记录所有脏页
	for _, d := range dirty {
		wal.RecordPage(d.PageID, d.Data)
	}

	// 4. HeaderPage 也纳入 (pageID=0)
	wal.RecordPage(0, marshalHeader(db.header))

	// 5. 原子提交
	return wal.Commit()
}

// ======================== 防抖自动提交 ========================

// SetAutoCommitDelay 设置防抖自动提交延迟
// delay=0 表示禁用自动提交 (仅手动 Commit)
func (db *DB) SetAutoCommitDelay(d time.Duration) {
	db.autoCommitMu.Lock()
	defer db.autoCommitMu.Unlock()
	db.autoCommitDelay = d
	if db.autoCommitTimer != nil {
		db.autoCommitTimer.Stop()
		db.autoCommitTimer = nil
	}
}

// NotifyWrite 写操作通知 (InsertRow/DeleteRow/UpdateRow 完成后调用)
// 重置防抖定时器: 如果延迟期间没有新操作, 则自动 Commit
func (db *DB) NotifyWrite() {
	if db.autoCommitDelay <= 0 {
		return
	}
	db.autoCommitMu.Lock()
	defer db.autoCommitMu.Unlock()

	// 复用同一个 timer 对象，避免频繁创建/销毁
	if db.autoCommitTimer == nil {
		// 首次创建
		db.autoCommitTimer = time.AfterFunc(db.autoCommitDelay, func() {
			_ = db.Commit()
		})
	} else {
		// 后续重置（无需重新分配对象，减少 GC 压力）
		db.autoCommitTimer.Reset(db.autoCommitDelay)
	}
}

// startAutoCommit 启动防抖自动提交
func (db *DB) startAutoCommit() {
	// 不立即启动定时器, 等待第一次 NotifyWrite
}

// stopAutoCommit 停止防抖自动提交
func (db *DB) stopAutoCommit() {
	db.autoCommitMu.Lock()
	defer db.autoCommitMu.Unlock()
	if db.autoCommitTimer != nil {
		db.autoCommitTimer.Stop()
		db.autoCommitTimer = nil
	}
}

// SetDirtyThreshold 设置脏页阈值（页数），超过此值自动触发 Commit
func (db *DB) SetDirtyThreshold(threshold int) {
	db.file.SetDirtyThreshold(threshold, func() {
		_ = db.Commit() // Commit 内部会自动停止防抖定时器
	})
}

// LockWrites 锁定写入（禁止 Insert/Update/Delete）
func (db *DB) LockWrites() {
	db.file.LockWrites()
}

// UnlockWrites 解锁写入（允许 Insert/Update/Delete）
func (db *DB) UnlockWrites() {
	db.file.UnlockWrites()
}

// SetBufferPoolCapacity 设置 BufferPool 容量（页数）
func (db *DB) SetBufferPoolCapacity(capacity int) {
	db.file.SetBufferPoolCapacity(capacity)
}

// marshalHeader 将 HeaderPage 序列化为字节切片 (填充到 PageSize)
func marshalHeader(h *HeaderPage) []byte {
	buf := make([]byte, PageSize)
	src := (*[PageSize4K]byte)(unsafe.Pointer(h))[:]
	copy(buf, src)
	return buf
}

// ======================== 调试诊断 ========================

// DebugDBState 打印数据库调试信息
func (db *DB) DebugDBState() {
	fmt.Printf("[DEBUG] Header: TablePage=%d NextFreePage=%d FreePageStart=%d FirstWAL=%d LastWAL=%d NextTxn=%d\n",
		db.header.TablePage, db.header.NextFreePage, db.header.FreePageStart,
		db.header.FirstWALPage, db.header.LastWALPage, db.header.NextTxnID)
	fmt.Printf("[DEBUG] BufferPool: total=%d dirty=%d\n", db.file.bp.Len(), db.file.bp.DirtyCount())
}

// DebugTableState 打印表调试信息 (从磁盘读取, 绕过 BufferPool)
func (db *DB) DebugTableState(tableName string) {
	tbl, ok := db.tables.Find(tableName)
	if !ok {
		fmt.Printf("[DEBUG] Table '%s' not found\n", tableName)
		return
	}
	fmt.Printf("[DEBUG] Table '%s': TotalRows=%d (memory)\n", tableName, tbl.info.TotalRows)

	// 从磁盘读取 TablePage 检查 TotalRows
	buf := make([]byte, PageSize)
	if err := db.file.ReadPage(tbl.TablePageID(), buf); err == nil {
		info := (*TableInfo)(unsafe.Pointer(&buf[tbl.TableStartOffset()]))
		fmt.Printf("[DEBUG] TablePage on disk: TotalRows=%d FieldCount=%d\n", info.TotalRows, info.FieldCount)
	}

	// 检查索引元数据
	im := tbl.IndexMeta()
	if im != nil {
		for i := range im.metas {
			m := &im.metas[i]
			if m.Flag&0x80 != 0 {
				continue
			}
			fmt.Printf("[DEBUG] Index '%s': EntryPage=%d IndexType=%d\n",
				string(m.IndexName[:]), m.EntryPage, m.IndexType)

			// 从磁盘读取 IndexMetaPage 检查 EntryPage
			imBuf := make([]byte, PageSize)
			if err := db.file.ReadPage(im.PageID(), imBuf); err == nil {
				// 读取磁盘上的 EntryPage
				diskMeta := (*IndexMeta)(unsafe.Pointer(&imBuf[im.HdrSize()]))
				fmt.Printf("[DEBUG] IndexMetaPage on disk: EntryPage=%d\n", diskMeta.EntryPage)
			}

			// 统计 B+ 树 DataPage 数量
			if m.IndexType == page.IndexTypePrimary {
				handler, _ := index.Compose.GetHandler(m.IndexType)
				if handler != nil {
					dpIDs := handler.GetDataPageIDs(tbl.handle, m.EntryPage)
					totalRows := 0
					for _, dpID := range dpIDs {
						dpBuf := make([]byte, PageSize)
						if err := db.file.ReadPage(dpID, dpBuf); err == nil {
							// 读取 DataPage 的行数 (在 offset 18 处的 4 字节 Count)
							count := binary.LittleEndian.Uint32(dpBuf[18:])
							totalRows += int(count)
						}
					}
					fmt.Printf("[DEBUG] B+ tree: %d DataPages, %d total rows\n", len(dpIDs), totalRows)
				}
			}
		}
	}
}

// DebugFreePageCount 返回空闲页数量 (调试用)
func (db *DB) DebugFreePageCount() int {
	return int(db.header.FreePageCount)
}