main.go

package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"os/signal"
	"syscall"
	"time"

	"pg2sqlite/internal/config"
	"pg2sqlite/internal/converter"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var (
	cfgFile string
	cfg     *config.Config
)

// versionCmd 版本命令
var versionCmd = &cobra.Command{
	Use:   "version",
	Short: "显示版本信息",
	Long:  "显示 pg2sqlite 的版本信息，包括版本号、构建时间和Git提交信息",
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Println(config.GetVersionInfo())
	},
}

// rootCmd 根命令
var rootCmd = &cobra.Command{
	Use:   "pg2sqlite",
	Short: "PostgreSQL 到 SQLite 转换工具",
	Long: `pg2sqlite 是一个高性能的 PostgreSQL 到 SQLite 转换工具。

支持功能:
- 批量数据转换
- 多表并行处理
- 自动数据类型转换
- 进度跟踪
- 错误重试机制
- 配置文件支持`,
	Run: runConvert,
}

// init 初始化命令
func init() {
	cobra.OnInitialize(initConfig)

	rootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "配置文件路径 (默认: config.yaml)")
	rootCmd.PersistentFlags().String("pg-host", "", "PostgreSQL 主机")
	rootCmd.PersistentFlags().Int("pg-port", 5432, "PostgreSQL 端口")
	rootCmd.PersistentFlags().String("pg-database", "", "PostgreSQL 数据库名")
	rootCmd.PersistentFlags().String("pg-username", "", "PostgreSQL 用户名")
	rootCmd.PersistentFlags().String("pg-password", "", "PostgreSQL 密码")
	rootCmd.PersistentFlags().String("sqlite-path", "", "SQLite 文件路径")
	rootCmd.PersistentFlags().Int64("test-limit", 0, "测试数据限制 (0表示无限制)")
	rootCmd.PersistentFlags().Int("batch-size", 1000, "批处理大小")
	rootCmd.PersistentFlags().Int("max-workers", 4, "最大工作协程数")
	rootCmd.PersistentFlags().Bool("drop-existing", false, "删除现有表")
	rootCmd.PersistentFlags().Bool("progress-log", true, "显示进度日志")

	// 绑定命令行参数到 viper
	viper.BindPFlag("database.postgresql.host", rootCmd.PersistentFlags().Lookup("pg-host"))
	viper.BindPFlag("database.postgresql.port", rootCmd.PersistentFlags().Lookup("pg-port"))
	viper.BindPFlag("database.postgresql.database", rootCmd.PersistentFlags().Lookup("pg-database"))
	viper.BindPFlag("database.postgresql.username", rootCmd.PersistentFlags().Lookup("pg-username"))
	viper.BindPFlag("database.postgresql.password", rootCmd.PersistentFlags().Lookup("pg-password"))
	viper.BindPFlag("database.sqlite.path", rootCmd.PersistentFlags().Lookup("sqlite-path"))
	viper.BindPFlag("conversion.test_limit", rootCmd.PersistentFlags().Lookup("test-limit"))
	viper.BindPFlag("conversion.batch_size", rootCmd.PersistentFlags().Lookup("batch-size"))
	viper.BindPFlag("conversion.max_workers", rootCmd.PersistentFlags().Lookup("max-workers"))
	viper.BindPFlag("conversion.drop_existing", rootCmd.PersistentFlags().Lookup("drop-existing"))
	viper.BindPFlag("conversion.progress_log", rootCmd.PersistentFlags().Lookup("progress-log"))
}

// initConfig 初始化配置
func initConfig() {
	// 设置默认值 (方案3)
	viper.SetDefault("conversion.batch_size", 1000)
	viper.SetDefault("conversion.max_workers", 4)
	viper.SetDefault("conversion.retry_count", 3)
	viper.SetDefault("conversion.retry_delay", time.Second)
	viper.SetDefault("database.sqlite.journal_mode", "WAL")
	viper.SetDefault("database.sqlite.performance.cache_size_mb", 500)
	viper.SetDefault("database.sqlite.performance.mmap_size_mb", 2048)
	viper.SetDefault("database.sqlite.performance.page_size_kb", 32)
	viper.SetDefault("database.sqlite.performance.busy_timeout_ms", 30000)
	viper.SetDefault("conversion.index.create_timing", "after")
	viper.SetDefault("conversion.index.batch_size", 1000)
	viper.SetDefault("database.postgresql.ssl_mode", "prefer")
	viper.SetDefault("database.postgresql.schema", "public")

	// 表发现配置默认值
	viper.SetDefault("conversion.table_discovery.mode", "manual")
	viper.SetDefault("conversion.table_discovery.exclude_tables", []string{})
	viper.SetDefault("conversion.table_discovery.include_patterns", []string{})
	viper.SetDefault("conversion.table_discovery.exclude_patterns", []string{})
	viper.SetDefault("conversion.table_discovery.min_rows", 0)
	viper.SetDefault("conversion.table_discovery.max_rows", 0)

	// 自动索引配置默认值
	viper.SetDefault("conversion.auto_index.enabled", false)
	viper.SetDefault("conversion.auto_index.primary_key_index", true)
	viper.SetDefault("conversion.auto_index.foreign_key_index", true)
	viper.SetDefault("conversion.auto_index.unique_index", true)
	viper.SetDefault("conversion.auto_index.skip_system_tables", true)

	if cfgFile != "" {
		viper.SetConfigFile(cfgFile)
	} else {
		viper.AddConfigPath(".")
		viper.AddConfigPath("./config")
		viper.SetConfigName("config")
		viper.SetConfigType("yaml")
	}

	viper.AutomaticEnv()

	if err := viper.ReadInConfig(); err != nil {
		if _, ok := err.(viper.ConfigFileNotFoundError); !ok {
			log.Fatalf("读取配置文件失败: %v", err)
		}
		log.Println("未找到配置文件，使用默认配置")
	} else {
		log.Printf("使用配置文件: %s", viper.ConfigFileUsed())
	}

	// 解析配置 - 使用直接解析方法
	cfg = &config.Config{}

	// 先尝试Unmarshal
	if err := viper.Unmarshal(cfg); err != nil {
		log.Printf("⚠️ Unmarshal失败，使用直接解析: %v", err)
	}

	// 使用直接解析方法确保配置正确
	cfg.Database.PostgreSQL.Host = viper.GetString("database.postgresql.host")
	cfg.Database.PostgreSQL.Port = viper.GetInt("database.postgresql.port")
	cfg.Database.PostgreSQL.Database = viper.GetString("database.postgresql.database")
	cfg.Database.PostgreSQL.Username = viper.GetString("database.postgresql.username")
	cfg.Database.PostgreSQL.Password = viper.GetString("database.postgresql.password")
	cfg.Database.PostgreSQL.SSLMode = viper.GetString("database.postgresql.ssl_mode")

	cfg.Database.SQLite.Path = viper.GetString("database.sqlite.path")
	cfg.Database.SQLite.JournalMode = viper.GetString("database.sqlite.journal_mode")

	cfg.Conversion.BatchSize = viper.GetInt("conversion.batch_size")
	cfg.Conversion.MaxWorkers = viper.GetInt("conversion.max_workers")
	cfg.Conversion.RetryCount = viper.GetInt("conversion.retry_count")
	cfg.Conversion.RetryDelay = viper.GetDuration("conversion.retry_delay")
	cfg.Conversion.ProgressLog = viper.GetBool("conversion.progress_log")
	cfg.Conversion.DropExisting = viper.GetBool("conversion.drop_existing")
	cfg.Conversion.SkipExisting = viper.GetBool("conversion.skip_existing")
	cfg.Conversion.BackupFile = viper.GetString("conversion.backup_file")
	cfg.Conversion.TestLimit = viper.GetInt64("conversion.test_limit")

	cfg.Conversion.Index.CreateTiming = viper.GetString("conversion.index.create_timing")
	cfg.Conversion.Index.BatchSize = viper.GetInt("conversion.index.batch_size")
	cfg.Conversion.Index.SkipUnsupported = viper.GetBool("conversion.index.skip_unsupported")
	cfg.Conversion.Index.OptimizeFor = viper.GetString("conversion.index.optimize_for")

	// 表发现配置
	cfg.Conversion.TableDiscovery.Mode = viper.GetString("conversion.table_discovery.mode")
	cfg.Conversion.TableDiscovery.ExcludeTables = viper.GetStringSlice("conversion.table_discovery.exclude_tables")
	cfg.Conversion.TableDiscovery.IncludePatterns = viper.GetStringSlice("conversion.table_discovery.include_patterns")
	cfg.Conversion.TableDiscovery.ExcludePatterns = viper.GetStringSlice("conversion.table_discovery.exclude_patterns")
	cfg.Conversion.TableDiscovery.MinRows = viper.GetInt64("conversion.table_discovery.min_rows")
	cfg.Conversion.TableDiscovery.MaxRows = viper.GetInt64("conversion.table_discovery.max_rows")

	// 自动索引配置
	cfg.Conversion.AutoIndex.Enabled = viper.GetBool("conversion.auto_index.enabled")
	cfg.Conversion.AutoIndex.PrimaryKeyIndex = viper.GetBool("conversion.auto_index.primary_key_index")
	cfg.Conversion.AutoIndex.ForeignKeyIndex = viper.GetBool("conversion.auto_index.foreign_key_index")
	cfg.Conversion.AutoIndex.UniqueIndex = viper.GetBool("conversion.auto_index.unique_index")
	cfg.Conversion.AutoIndex.SkipSystemTables = viper.GetBool("conversion.auto_index.skip_system_tables")

	log.Printf("✅ 直接解析完成")

	// 配置解析完成

	// 验证配置
	if err := cfg.Validate(); err != nil {
		log.Fatalf("配置验证失败: %v", err)
	}

	// 配置解析完成，无需手动设置
}

// runConvert 运行转换
func runConvert(cmd *cobra.Command, args []string) {
	log.Printf("🚀 %s", config.GetVersionInfo())
	log.Printf("🚀 pg2sqlite 转换工具启动")
	log.Printf("📋 配置信息:")
	log.Printf("   - PostgreSQL: %s:%d/%s",
		cfg.Database.PostgreSQL.Host,
		cfg.Database.PostgreSQL.Port,
		cfg.Database.PostgreSQL.Database)
	log.Printf("   - SQLite: %s", cfg.Database.SQLite.Path)
	log.Printf("   - 批处理大小: %d", cfg.Conversion.BatchSize)
	log.Printf("   - 最大工作协程: %d", cfg.Conversion.MaxWorkers)
	if cfg.Conversion.TestLimit > 0 {
		log.Printf("   - 测试限制: %d 条记录", cfg.Conversion.TestLimit)
	}

	// 创建转换器
	converter, err := converter.NewConverter(cfg)
	if err != nil {
		log.Fatalf("创建转换器失败: %v", err)
	}
	defer converter.Close()

	// 设置信号处理
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		<-sigChan
		log.Println("\n🛑 收到中断信号，正在安全退出...")
		cancel()
	}()

	// 执行转换
	startTime := time.Now()
	if err := converter.Convert(ctx); err != nil {
		log.Fatalf("转换失败: %v", err)
	}

	duration := time.Since(startTime)
	log.Printf("🎉 转换完成，总耗时: %v", duration)
}

// main 主函数
func main() {
	// 添加版本命令
	rootCmd.AddCommand(versionCmd)

	// 检查是否是版本命令
	if len(os.Args) > 1 && os.Args[1] == "version" {
		fmt.Println(config.GetVersionInfo())
		return
	}

	if err := rootCmd.Execute(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}