package main

import (
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"strings"
	"time"

	"github.com/go-sql-driver/mysql"

	"github.com/tidbcloud/AgentTeam/envcompat"
	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/migrations"
	"github.com/tidbcloud/AgentTeam/internal/schema"
	_ "github.com/tidbcloud/AgentTeam/internal/store/organization"
	"github.com/tidbcloud/AgentTeam/internal/tidb"
)

const defaultConnectTimeout = 5 * time.Second

type migrationReportingAdmin interface {
	ApplyMigrationsWithReport(context.Context, coredb.AdminDB) (migrations.ApplyReport, error)
}

func main() {
	ctx := context.Background()

	command := "init"
	if len(os.Args) > 1 {
		command = strings.TrimSpace(os.Args[1])
	}

	switch command {
	case "init":
		if err := runInit(ctx); err != nil {
			fmt.Fprintf(os.Stderr, "agentteam-schema init failed: %v\n", err)
			os.Exit(1)
		}
	case "verify":
		if err := runVerify(ctx); err != nil {
			fmt.Fprintf(os.Stderr, "agentteam-schema verify failed: %v\n", err)
			os.Exit(1)
		}
	case "-h", "--help", "help":
		printUsage()
	default:
		printUsage()
		fmt.Fprintf(os.Stderr, "\nunknown command %q\n", command)
		os.Exit(1)
	}
}

func runInit(ctx context.Context) error {
	meshXEnabled, err := meshXSchemaEnabled()
	if err != nil {
		return err
	}
	if meshXEnabled {
		// Validate every enabled runtime-owned schema setting before database
		// creation, startup preflight, or migrations can produce side effects.
		if err := schema.ValidateConfigFromEnv(); err != nil {
			return err
		}
		if err := ensureAllRuntimeDatabases(ctx); err != nil {
			return err
		}
	}
	admin, db, err := openCurrentAdminDBForInit(ctx)
	if err != nil {
		return err
	}
	defer db.Close()

	fmt.Println("agentteam-schema init: database=tidb")
	fmt.Println("running startup preflight")
	if err := admin.RunStartupPreflight(ctx, db); err != nil {
		return err
	}
	fmt.Println("startup preflight ok")

	report, err := applyMigrationsWithReport(ctx, admin, db)
	if err != nil {
		return err
	}
	writeApplyReport(os.Stdout, report)

	fmt.Println("verifying current schema")
	if err := admin.RequireSchemaCurrent(ctx, db); err != nil {
		return err
	}
	fmt.Println("schema current")
	fmt.Println("schema initialized")
	if !meshXEnabled {
		return nil
	}
	return schema.MigrateFromEnv(ctx)
}

func meshXSchemaEnabled() (bool, error) {
	return envcompat.Bool("AGENTTEAM_MESHX_ENABLED", false)
}

func ensureAllRuntimeDatabases(ctx context.Context) error {
	agentteamDSN, _, _ := envcompat.Lookup("AGENTTEAM_DATABASE_DSN", "AGENTTEAM_CORE_STORAGE_DSN")
	dsns := []struct {
		name string
		dsn  string
	}{
		{name: "AgentTeam", dsn: agentteamDSN},
		{name: "AGS", dsn: os.Getenv("MESHX_AGS_DATABASE_DSN")},
		{name: "MeshX", dsn: os.Getenv("MESHX_DATABASE_DSN")},
	}
	connectTimeout, err := parseConnectTimeout()
	if err != nil {
		return err
	}
	for _, target := range dsns {
		dsn := strings.TrimSpace(target.dsn)
		if dsn == "" {
			return fmt.Errorf("%s database DSN is required", target.name)
		}
		lowerDSN := strings.ToLower(dsn)
		if lowerDSN == ":memory:" || strings.HasPrefix(lowerDSN, "file:") || strings.HasPrefix(lowerDSN, "sqlite:") {
			continue
		}
		if err := ensureTiDBDatabaseExists(ctx, dsn, connectTimeout); err != nil {
			return fmt.Errorf("ensure %s database: %w", target.name, err)
		}
	}
	return nil
}

func runVerify(ctx context.Context) error {
	meshXEnabled, err := meshXSchemaEnabled()
	if err != nil {
		return err
	}
	admin, db, err := openCurrentAdminDB(ctx)
	if err != nil {
		return err
	}
	defer db.Close()

	fmt.Println("agentteam-schema verify: database=tidb")
	fmt.Println("running startup preflight")
	if err := admin.RunStartupPreflight(ctx, db); err != nil {
		return err
	}
	fmt.Println("startup preflight ok")
	fmt.Println("verifying current schema")
	if err := admin.RequireSchemaCurrent(ctx, db); err != nil {
		return err
	}
	if state, err := admin.InspectMigrationState(ctx, db); err == nil {
		writeVersionState(os.Stdout, "migration state", migrations.VersionState{
			Known:   true,
			Current: state.CurrentVersion,
			Target:  state.TargetVersion,
			Pending: state.Pending,
		})
	}
	fmt.Println("schema current")
	fmt.Println("schema verified")
	if !meshXEnabled {
		return nil
	}
	return schema.VerifyFromEnv(ctx)
}

func applyMigrationsWithReport(ctx context.Context, admin coredb.Admin, db coredb.AdminDB) (migrations.ApplyReport, error) {
	if reportingAdmin, ok := admin.(migrationReportingAdmin); ok {
		report, err := reportingAdmin.ApplyMigrationsWithReport(ctx, db)
		if report.Database == "" {
			report.Database = admin.Database()
		}
		return report, err
	}
	err := admin.ApplyMigrations(ctx, db)
	return migrations.ApplyReport{Database: admin.Database()}, err
}

func writeApplyReport(w io.Writer, report migrations.ApplyReport) {
	if report.Database != "" {
		fmt.Fprintf(w, "migration database: %s\n", report.Database)
	}
	if report.Before.Known {
		writeVersionState(w, "migration state before init", report.Before)
	}
	if report.After.Known {
		writeVersionState(w, "migration state after init", report.After)
	}
	fmt.Fprintln(w, "executed migration sql files:")
	if len(report.ExecutedFiles) == 0 {
		fmt.Fprintln(w, "- <none>")
	} else {
		for _, file := range report.ExecutedFiles {
			if file.Version > 0 {
				fmt.Fprintf(w, "- %s (version=%d)\n", file.Path, file.Version)
				continue
			}
			fmt.Fprintf(w, "- %s\n", file.Path)
		}
	}
	if len(report.Notes) > 0 {
		fmt.Fprintln(w, "migration notes:")
		for _, note := range report.Notes {
			fmt.Fprintf(w, "- %s\n", note)
		}
	}
}

func writeVersionState(w io.Writer, label string, state migrations.VersionState) {
	if !state.Known {
		return
	}
	fmt.Fprintf(w, "%s: current=%d target=%d pending=%t\n", label, state.Current, state.Target, state.Pending)
}

func openCurrentAdminDB(ctx context.Context) (coredb.Admin, coredb.AdminDB, error) {
	admin, databaseURL, connectTimeout, err := currentAdminConfig()
	if err != nil {
		return nil, nil, err
	}

	db, err := openAdminDB(ctx, admin, databaseURL, connectTimeout)
	if err != nil {
		return nil, nil, fmt.Errorf("connect core main-path database: %w", err)
	}
	return admin, db, nil
}

func openCurrentAdminDBForInit(ctx context.Context) (coredb.Admin, coredb.AdminDB, error) {
	admin, databaseURL, connectTimeout, err := currentAdminConfig()
	if err != nil {
		return nil, nil, err
	}

	return openAdminDBForInit(ctx, admin, databaseURL, connectTimeout, openAdminDB, ensureTiDBDatabaseExists)
}

func currentAdminConfig() (coredb.Admin, string, time.Duration, error) {
	databaseURL, err := currentDatabaseURL()
	if err != nil {
		return nil, "", 0, err
	}

	connectTimeout, err := parseConnectTimeout()
	if err != nil {
		return nil, "", 0, err
	}
	return coredb.NewAdmin(), databaseURL, connectTimeout, nil
}

func openAdminDBForInit(
	ctx context.Context,
	admin coredb.Admin,
	databaseURL string,
	connectTimeout time.Duration,
	open func(context.Context, coredb.Admin, string, time.Duration) (coredb.AdminDB, error),
	ensureTiDBDatabase func(context.Context, string, time.Duration) error,
) (coredb.Admin, coredb.AdminDB, error) {
	db, err := open(ctx, admin, databaseURL, connectTimeout)
	if err == nil {
		return admin, db, nil
	}
	if !isUnknownTiDBDatabaseError(err) {
		return nil, nil, fmt.Errorf("connect core main-path database: %w", err)
	}
	if err := ensureTiDBDatabase(ctx, databaseURL, connectTimeout); err != nil {
		return nil, nil, err
	}

	db, err = open(ctx, admin, databaseURL, connectTimeout)
	if err != nil {
		return nil, nil, fmt.Errorf("connect core main-path database: %w", err)
	}
	return admin, db, nil
}

func openAdminDB(ctx context.Context, admin coredb.Admin, databaseURL string, connectTimeout time.Duration) (coredb.AdminDB, error) {
	return admin.Open(ctx, databaseURL, connectTimeout)
}

func ensureTiDBDatabaseExists(ctx context.Context, databaseURL string, connectTimeout time.Duration) error {
	cfg, err := mysql.ParseDSN(strings.TrimSpace(databaseURL))
	if err != nil {
		return fmt.Errorf("parse tidb dsn: %w", err)
	}
	databaseName := strings.TrimSpace(cfg.DBName)
	if databaseName == "" {
		return fmt.Errorf("tidb dsn must include a database name")
	}

	adminCfg := *cfg
	adminCfg.DBName = ""
	adminDB, err := tidb.NewAdmin().Open(ctx, adminCfg.FormatDSN(), connectTimeout)
	if err != nil {
		return fmt.Errorf("open tidb admin connection: %w", err)
	}
	defer func() { _ = adminDB.Close() }()

	if _, err := adminDB.ExecContext(ctx, "CREATE DATABASE IF NOT EXISTS "+quoteTiDBIdentifier(databaseName)); err != nil {
		return fmt.Errorf("create tidb database %s: %w", databaseName, err)
	}
	return nil
}

func isUnknownTiDBDatabaseError(err error) bool {
	var mysqlErr *mysql.MySQLError
	return errors.As(err, &mysqlErr) && mysqlErr.Number == 1049
}

func quoteTiDBIdentifier(identifier string) string {
	return "`" + strings.ReplaceAll(identifier, "`", "``") + "`"
}

func currentDatabaseURL() (string, error) {
	databaseURL, _, ok := envcompat.Lookup("AGENTTEAM_DATABASE_DSN", "AGENTTEAM_CORE_STORAGE_DSN")
	if ok {
		return databaseURL, nil
	}
	return "", fmt.Errorf("AGENTTEAM_DATABASE_DSN is required")
}

func parseConnectTimeout() (time.Duration, error) {
	raw := strings.TrimSpace(os.Getenv("AGENTTEAM_DATABASE_CONNECT_TIMEOUT"))
	if raw == "" {
		return defaultConnectTimeout, nil
	}

	duration, err := time.ParseDuration(raw)
	if err != nil {
		return 0, fmt.Errorf("parse AGENTTEAM_DATABASE_CONNECT_TIMEOUT: %w", err)
	}
	if duration <= 0 {
		return 0, fmt.Errorf("AGENTTEAM_DATABASE_CONNECT_TIMEOUT must be positive")
	}
	return duration, nil
}

func printUsage() {
	fmt.Fprintf(os.Stderr, "Usage: %s [init|verify]\n", os.Args[0])
	fmt.Fprintln(os.Stderr)
	fmt.Fprintln(os.Stderr, "Commands:")
	fmt.Fprintln(os.Stderr, "  init    initialize the latest runtime schema on an empty database and verify it")
	fmt.Fprintln(os.Stderr, "  verify  require that the current schema matches the running code")
	fmt.Fprintln(os.Stderr)
	fmt.Fprintln(os.Stderr, "MeshX and AGS schemas are included only when AGENTTEAM_MESHX_ENABLED=true.")
}