package main

import (
	"database/sql"
	"log"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

var db *sql.DB

func InitDB() {
	var err error
	db, err = sql.Open("sqlite3", "./chat.db")
	if err != nil {
		log.Fatalf("打开数据库失败: %v", err)
	}

	createTableSQL := `
    CREATE TABLE IF NOT EXISTS messages (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        sender TEXT NOT NULL,
        receiver TEXT NOT NULL,
        content TEXT NOT NULL,
        timestamp DATETIME NOT NULL
    );
    `
	_, err = db.Exec(createTableSQL)
	if err != nil {
		log.Fatalf("创建表失败: %v", err)
	}
}

func SaveMessage(sender, receiver, content string) error {
	stmt, err := db.Prepare("INSERT INTO messages(sender, receiver, content, timestamp) VALUES (?, ?, ?, ?)")
	if err != nil {
		return err
	}
	defer stmt.Close()

	_, err = stmt.Exec(sender, receiver, content, time.Now())
	return err
}

func GetMessages(sender, receiver string) ([]Message, error) {
	rows, err := db.Query("SELECT sender, receiver, content, timestamp FROM messages WHERE (sender = ? AND receiver = ?) OR (sender = ? AND receiver = ?) ORDER BY timestamp ASC", sender, receiver, receiver, sender)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var messages []Message
	for rows.Next() {
		var msg Message
		var ts string
		err = rows.Scan(&msg.Sender, &msg.Receiver, &msg.Content, &ts)
		if err != nil {
			return nil, err
		}
		msg.Timestamp, _ = time.Parse(time.RFC3339, ts)
		messages = append(messages, msg)
	}
	return messages, nil
}

type Message struct {
	Sender    string
	Receiver  string
	Content   string
	Timestamp time.Time
}