package migrate

import (
	sdkmigrate "github.com/tidbcloud/aisupport/packages/sdk-go/infra/database/migrate"
	"github.com/tidbcloud/aisupport/services/agent-gateway/models"
	"gorm.io/gorm"
)

var Migrator = sdkmigrate.NewMigrator(
	sdkmigrate.WithTable("agent_gateway_migrations"),
	sdkmigrate.WithInitDBFunc(initDB),
)

var _ = Migrator.Add("202608311200_create_agent_gateway_tables", initDB)
var _ = Migrator.Add("202609021200_rename_agent_provider_to_agent_id", renameAgentProviderToAgentID)
var _ = Migrator.Add("202609071200_add_task_agent_model", addTaskAgentModel)

func initDB(db *gorm.DB) error {
	return db.AutoMigrate(&models.Session{}, &models.LarkSessionContext{}, &models.Task{})
}

func renameAgentProviderToAgentID(db *gorm.DB) error {
	migrator := db.Migrator()
	if migrator.HasColumn(&models.Session{}, "agent_id") {
		return nil
	}
	return migrator.RenameColumn(&models.Session{}, "agent_provider", "agent_id")
}

func addTaskAgentModel(db *gorm.DB) error {
	migrator := db.Migrator()
	if migrator.HasColumn(&models.Task{}, "agent_model") {
		return nil
	}
	return migrator.AddColumn(&models.Task{}, "AgentModel")
}