package migrate

import (
	"strings"
	"testing"

	"github.com/stretchr/testify/require"
	"gorm.io/driver/mysql"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"

	"github.com/tidbcloud/aisupport/services/agent-gateway/models"
)

func TestMigratorCreatesAgentGatewaySchema(t *testing.T) {
	t.Parallel()

	db, err := gorm.Open(sqlite.Open(t.TempDir()+"/agent-gateway.db"), &gorm.Config{})
	require.NoError(t, err)
	require.NoError(t, Migrator.Run(db))

	require.True(t, db.Migrator().HasTable(&models.Session{}))
	require.True(t, db.Migrator().HasTable(&models.LarkSessionContext{}))
	require.True(t, db.Migrator().HasTable(&models.Task{}))
	require.True(t, db.Migrator().HasTable("agent_gateway_migrations"))
	require.True(t, db.Migrator().HasColumn(&models.Session{}, "agent_id"))
	require.False(t, db.Migrator().HasColumn("agent_gateway_sessions", "agent_provider"))
	require.True(t, db.Migrator().HasColumn(&models.Session{}, "source_initiator_email"))
	require.False(t, db.Migrator().HasColumn(&models.Session{}, "source_initiator_id"))
	require.True(t, db.Migrator().HasColumn(&models.Task{}, "source_sender_email"))
	require.True(t, db.Migrator().HasColumn(&models.Task{}, "working_directory"))
	require.False(t, db.Migrator().HasColumn(&models.Task{}, "source_sender_id"))
	require.False(t, db.Migrator().HasColumn(&models.Session{}, "lark_chat_id"))
	require.False(t, db.Migrator().HasColumn(&models.Session{}, "lark_thread_id"))
	require.False(t, db.Migrator().HasColumn(&models.Session{}, "lark_anchor_message_id"))
	require.False(t, db.Migrator().HasColumn(&models.Session{}, "lark_entry_message_id"))
	require.True(t, db.Migrator().HasColumn(&models.LarkSessionContext{}, "chat_id"))
	require.True(t, db.Migrator().HasColumn(&models.LarkSessionContext{}, "thread_id"))
	require.True(t, db.Migrator().HasColumn(&models.LarkSessionContext{}, "anchor_message_id"))
	require.True(t, db.Migrator().HasColumn(&models.LarkSessionContext{}, "entry_message_id"))
	require.True(t, db.Migrator().HasIndex(&models.LarkSessionContext{}, "ChatID"))
	require.True(t, db.Migrator().HasIndex(&models.LarkSessionContext{}, "ThreadID"))
	require.True(t, db.Migrator().HasIndex(&models.LarkSessionContext{}, "AnchorMessageID"))
	require.True(t, db.Migrator().HasIndex(&models.LarkSessionContext{}, "EntryMessageID"))
	require.False(t, db.Migrator().HasTable("sessions"))
	require.False(t, db.Migrator().HasTable("tasks"))

	requireColumnTypes(t, db, &models.Session{}, map[string]string{
		"session_id":             "varchar(255)",
		"agent_id":               "varchar(64)",
		"source_initiator_email": "varchar(320)",
		"created_at":             "datetime",
	})
	requireColumnTypes(t, db, &models.LarkSessionContext{}, map[string]string{
		"session_id": "varchar(255)",
		"chat_id":    "varchar(255)",
	})
	requireColumnTypes(t, db, &models.Task{}, map[string]string{
		"task_id":             "varchar(255)",
		"status":              "varchar(64)",
		"source_sender_email": "varchar(320)",
		"prompt":              "text",
		"working_directory":   "text",
		"agent_model":         "varchar(255)",
		"error_message":       "text",
		"updated_at":          "datetime",
	})
}

func TestAgentGatewayModelMySQLColumnTypes(t *testing.T) {
	t.Parallel()

	db, err := gorm.Open(mysql.New(mysql.Config{SkipInitializeWithVersion: true}), &gorm.Config{
		DisableAutomaticPing: true,
	})
	require.NoError(t, err)

	requireMySQLColumnTypes(t, db, &models.Session{}, map[string]string{
		"session_id":             "varchar(255)",
		"agent_id":               "varchar(64)",
		"agent_session_id":       "varchar(255)",
		"active_turn_id":         "varchar(255)",
		"source_type":            "varchar(64)",
		"source_instance_id":     "varchar(255)",
		"source_session_key":     "varchar(255)",
		"source_initiator_email": "varchar(320)",
		"status":                 "varchar(64)",
		"created_at":             "datetime(3)",
		"updated_at":             "datetime(3)",
	})
	requireMySQLColumnTypes(t, db, &models.LarkSessionContext{}, map[string]string{
		"session_id":        "varchar(255)",
		"chat_id":           "varchar(255)",
		"thread_id":         "varchar(255)",
		"anchor_message_id": "varchar(255)",
		"entry_message_id":  "varchar(255)",
	})
	requireMySQLColumnTypes(t, db, &models.Task{}, map[string]string{
		"task_id":             "varchar(255)",
		"session_id":          "varchar(255)",
		"agent_turn_id":       "varchar(255)",
		"source_message_id":   "varchar(255)",
		"source_sender_email": "varchar(320)",
		"status":              "varchar(64)",
		"prompt":              "text",
		"working_directory":   "text",
		"agent_model":         "varchar(255)",
		"error_message":       "text",
		"created_at":          "datetime(3)",
		"updated_at":          "datetime(3)",
	})

	// Include the default clause: TEXT alone is valid MySQL, but TEXT DEFAULT '' is not.
	stmt := &gorm.Statement{DB: db}
	require.NoError(t, stmt.Parse(&models.Task{}))
	field := stmt.Schema.LookUpField("agent_model")
	require.NotNil(t, field)
	require.Equal(t, "varchar(255) NOT NULL DEFAULT ''", db.Migrator().FullDataTypeOf(field).SQL)
}

func TestMigratorRenamesAgentProviderAndPreservesSessionData(t *testing.T) {
	t.Parallel()

	db, err := gorm.Open(sqlite.Open(t.TempDir()+"/agent-gateway-old.db"), &gorm.Config{})
	require.NoError(t, err)
	require.NoError(t, db.Exec(`CREATE TABLE agent_gateway_sessions (
session_id varchar(255) PRIMARY KEY,
agent_provider varchar(64) NOT NULL
)`).Error)
	require.NoError(t, db.Exec(`INSERT INTO agent_gateway_sessions (session_id, agent_provider) VALUES (?, ?)`, "session-1", "codex").Error)
	require.NoError(t, db.Exec(`CREATE TABLE agent_gateway_tasks (task_id varchar(255) PRIMARY KEY)`).Error)
	require.NoError(t, db.Exec(`CREATE TABLE agent_gateway_migrations (id varchar(255) PRIMARY KEY)`).Error)
	require.NoError(t, db.Exec(`INSERT INTO agent_gateway_migrations (id) VALUES (?)`, "202608311200_create_agent_gateway_tables").Error)

	require.NoError(t, Migrator.Run(db))
	require.True(t, db.Migrator().HasColumn("agent_gateway_sessions", "agent_id"))
	require.False(t, db.Migrator().HasColumn("agent_gateway_sessions", "agent_provider"))

	var agentID string
	require.NoError(t, db.Raw(`SELECT agent_id FROM agent_gateway_sessions WHERE session_id = ?`, "session-1").Scan(&agentID).Error)
	require.Equal(t, "codex", agentID)
}

func requireColumnTypes(t *testing.T, db *gorm.DB, model any, expected map[string]string) {
	t.Helper()

	columnTypes, err := db.Migrator().ColumnTypes(model)
	require.NoError(t, err)

	actual := make(map[string]string, len(columnTypes))
	for _, columnType := range columnTypes {
		databaseType, ok := columnType.ColumnType()
		require.True(t, ok, columnType.Name())
		actual[columnType.Name()] = strings.ToLower(databaseType)
	}
	for column, databaseType := range expected {
		require.Equal(t, databaseType, actual[column], column)
	}
}

func requireMySQLColumnTypes(t *testing.T, db *gorm.DB, model any, expected map[string]string) {
	t.Helper()

	stmt := &gorm.Statement{DB: db}
	require.NoError(t, stmt.Parse(model))
	for column, databaseType := range expected {
		field := stmt.Schema.LookUpField(column)
		require.NotNil(t, field, column)
		require.Equal(t, databaseType, strings.ToLower(db.Dialector.DataTypeOf(field)), column)
	}
}

func TestAddTaskAgentModelMigratesExistingRows(t *testing.T) {
	t.Parallel()
	db, err := gorm.Open(sqlite.Open(t.TempDir()+"/old.db"), &gorm.Config{})
	require.NoError(t, err)
	require.NoError(t, db.Exec("CREATE TABLE agent_gateway_tasks (task_id varchar(255) PRIMARY KEY, status varchar(64), updated_at datetime)").Error)
	require.NoError(t, db.Exec("INSERT INTO agent_gateway_tasks (task_id, status) VALUES ('old-task', 'RECEIVED')").Error)
	require.NoError(t, addTaskAgentModel(db))
	require.NoError(t, addTaskAgentModel(db))
	var task models.Task
	require.NoError(t, db.First(&task, "task_id = ?", "old-task").Error)
	require.Empty(t, task.AgentModel)
	require.Equal(t, "RECEIVED", task.Status)
	require.NoError(t, db.Model(&task).Update("agent_model", "requested-model").Error)
	require.NoError(t, db.First(&task, "task_id = ?", "old-task").Error)
	require.Equal(t, "requested-model", task.AgentModel)
}