/*
Copyright 2025-2026 linux.do
Modified by Arctel.net, 2026

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package db

import (
	"context"
	"log"
	"net"
	"net/url"
	"strconv"
	"time"

	"github.com/Rain-kl/Wavelet/internal/config"
	"github.com/glebarez/sqlite"
	"go.opentelemetry.io/otel/attribute"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	"gorm.io/plugin/dbresolver"
	"gorm.io/plugin/opentelemetry/tracing"
)

var (
	db *gorm.DB
)

func init() {
	if !config.Config.Database.Enabled {
		// PostgreSQL 禁用，使用 SQLite
		initSQLite()
		return
	}

	initPostgres()
}

// initSQLite 初始化 SQLite 数据库（PostgreSQL 禁用时的后备方案）
func initSQLite() {
	sqlitePath := config.Config.Database.SQLitePath
	if sqlitePath == "" {
		sqlitePath = "./data/wavelet.db"
	}

	var err error
	db, err = gorm.Open(sqlite.Open(sqlitePath), &gorm.Config{
		DisableForeignKeyConstraintWhenMigrating: true,
		Logger: &gormZapLogger{
			logLevel:                  parseLogLevel(config.Config.Database.LogLevel),
			slowThreshold:             config.Config.Database.SlowThreshold,
			ignoreRecordNotFoundError: config.Config.App.IsProduction(),
		},
	})
	if err != nil {
		log.Fatalf("[SQLite] init connection failed: %v\n", err)
	}

	log.Printf("[SQLite] initialized (path: %s)\n", sqlitePath)
}

// initPostgres 初始化 PostgreSQL 数据库
func initPostgres() {
	var err error
	dbConfig := config.Config.Database

	// 构建主库 DSN 并连接
	primaryDSN := buildDSN(dbConfig.Host, dbConfig.Port, dbConfig.Username, dbConfig.Password)

	pgConfig := postgres.Config{
		DSN:                  primaryDSN,
		PreferSimpleProtocol: dbConfig.PreferSimpleProtocol,
	}

	db, err = gorm.Open(postgres.New(pgConfig), &gorm.Config{
		DisableForeignKeyConstraintWhenMigrating: true,
		Logger: &gormZapLogger{
			logLevel:                  parseLogLevel(config.Config.Database.LogLevel),
			slowThreshold:             config.Config.Database.SlowThreshold,
			ignoreRecordNotFoundError: config.Config.App.IsProduction(),
		},
	})
	if err != nil {
		log.Fatalf("[PostgreSQL] init connection failed: %v\n", err)
	}

	// Trace 注入
	if err = db.Use(
		tracing.NewPlugin(
			tracing.WithoutMetrics(),
			tracing.WithAttributes(
				attribute.String("db.instance", dbConfig.Database),
				attribute.String("db.ip", dbConfig.Host),
				attribute.String("server.address", net.JoinHostPort(dbConfig.Host, strconv.Itoa(dbConfig.Port))),
				attribute.String("db.system", "PostgreSQL"),
			),
		),
	); err != nil {
		log.Fatalf("[PostgreSQL] init trace failed: %v\n", err)
	}

	if len(dbConfig.Replicas) > 0 {
		var replicaDialectors []gorm.Dialector
		for _, replica := range dbConfig.Replicas {
			username := replica.Username
			if username == "" {
				username = dbConfig.Username
			}
			password := replica.Password
			if password == "" {
				password = dbConfig.Password
			}
			replicaDSN := buildDSN(replica.Host, replica.Port, username, password)
			replicaDialectors = append(replicaDialectors, postgres.New(postgres.Config{
				DSN:                  replicaDSN,
				PreferSimpleProtocol: dbConfig.PreferSimpleProtocol,
			}))
		}

		resolver := dbresolver.Register(dbresolver.Config{
			Replicas: replicaDialectors,
			Policy:   dbresolver.RandomPolicy{},
		})

		resolver.SetMaxIdleConns(dbConfig.MaxIdleConn).
			SetMaxOpenConns(dbConfig.MaxOpenConn).
			SetConnMaxLifetime(time.Duration(dbConfig.ConnMaxLifetime) * time.Second).
			SetConnMaxIdleTime(time.Duration(dbConfig.ConnMaxIdleTime) * time.Second)

		if err = db.Use(resolver); err != nil {
			log.Fatalf("[PostgreSQL] init dbresolver failed: %v\n", err)
		}
		log.Printf("[PostgreSQL] initialized in Primary-Replica mode (%d replicas)\n", len(dbConfig.Replicas))
	} else {
		log.Println("[PostgreSQL] initialized in Standalone mode")
	}

	// 获取通用数据库对象设置连接池
	sqlDB, err := db.DB()
	if err != nil {
		log.Fatalf("[PostgreSQL] load sql db failed: %v\n", err)
	}

	sqlDB.SetMaxIdleConns(dbConfig.MaxIdleConn)
	sqlDB.SetMaxOpenConns(dbConfig.MaxOpenConn)
	sqlDB.SetConnMaxLifetime(time.Duration(dbConfig.ConnMaxLifetime) * time.Second)
	sqlDB.SetConnMaxIdleTime(time.Duration(dbConfig.ConnMaxIdleTime) * time.Second)

}

// buildDSN 构建 PostgreSQL DSN
func buildDSN(host string, port int, username, password string) string {
	cfg := config.Config.Database
	pqURL := &url.URL{
		Scheme: "postgres",
		Host:   net.JoinHostPort(host, strconv.Itoa(port)),
		Path:   cfg.Database,
	}
	if username != "" {
		pqURL.User = url.UserPassword(username, password)
	}

	query := pqURL.Query()
	sslMode := cfg.SSLMode
	if sslMode == "" {
		sslMode = "disable"
	}
	query.Set("sslmode", sslMode)
	if cfg.ApplicationName != "" {
		query.Set("application_name", cfg.ApplicationName)
	}
	if cfg.SearchPath != "" {
		query.Set("search_path", cfg.SearchPath)
	}
	if cfg.DefaultQueryExecMode != "" {
		query.Set("default_query_exec_mode", cfg.DefaultQueryExecMode)
	}
	if cfg.StatementCacheCapacity > 0 {
		query.Set("statement_cache_capacity", strconv.Itoa(cfg.StatementCacheCapacity))
	}

	rawQuery := query.Encode()
	if cfg.TimeZone != "" {
		if rawQuery != "" {
			rawQuery += "&"
		}
		rawQuery += "TimeZone=" + cfg.TimeZone
	}
	pqURL.RawQuery = rawQuery

	return pqURL.String()
}

// DB 返回带上下文追踪的 GORM 数据库实例
func DB(ctx context.Context) *gorm.DB {
	if db == nil {
		return nil
	}
	return db.WithContext(ctx)
}

// SetDB sets the package-level database instance for testing.
func SetDB(d *gorm.DB) {
	db = d
}