/*
 * Copyright 2021 PingCAP, Inc.
 *
 * 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,
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
	"context"
	"errors"
	"fmt"
	"log"
	"runtime"
	"strings"
	"time"

	"github.com/go-sql-driver/mysql"
	"golang.org/x/sync/errgroup"
	gormysql "gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

const (
	MaxDialTimeout   = 30000 // millisecond
	MaxReadTimeout   = 30000 // millisecond
	MaxWriteTimeout  = 30000 // millisecond
	MaxOpenConn      = 128
	MaxIdleConn      = 16
	MaxLifecycleConn = 300 // in second
	PWDPlaceHolder   = "__PWD__"
)

// MysqlConfig defines config for gorm dialer
type MysqlConfig struct {
	NeedMergePWD bool          `yaml:"need_merge_pwd"`
	DSN          string        `yaml:"dsn"`
	DialTimeout  time.Duration `yaml:"dial_timeout"`
	ReadTimeout  time.Duration `yaml:"read_timeout"`
	WriteTimeout time.Duration `yaml:"write_timeout"`
	MaxOpenConns int           `yaml:"max_open_conns"`
	MaxIdleConns int           `yaml:"max_idle_conns"`
	MaxLifeConns int           `yaml:"max_life_conns"`
}

func (c *MysqlConfig) MergePWD(pwd string) {
	c.DSN = strings.ReplaceAll(c.DSN, PWDPlaceHolder, pwd)
}

// FillWithDefaults apply default values for field with invalid value.
func (c *MysqlConfig) FillWithDefaults() {
	if c == nil {
		return
	}
	maxCPU := runtime.NumCPU()

	if c.DialTimeout <= 0 || c.DialTimeout > time.Duration(MaxDialTimeout*maxCPU) {
		c.DialTimeout = MaxDialTimeout
	}

	if c.ReadTimeout <= 0 || c.ReadTimeout > time.Duration(MaxReadTimeout*maxCPU) {
		c.ReadTimeout = MaxReadTimeout
	}

	if c.WriteTimeout <= 0 || c.WriteTimeout > time.Duration(MaxWriteTimeout*maxCPU) {
		c.WriteTimeout = MaxWriteTimeout
	}

	if c.MaxOpenConns <= 0 || c.MaxOpenConns > MaxOpenConn*maxCPU {
		c.MaxOpenConns = MaxOpenConn
	}

	if c.MaxIdleConns <= 0 || c.MaxIdleConns > MaxIdleConn*maxCPU {
		c.MaxIdleConns = MaxIdleConn
	}

	if c.MaxLifeConns <= 0 || c.MaxLifeConns > MaxLifecycleConn*maxCPU {
		c.MaxLifeConns = MaxLifecycleConn
	}
}

func (c *MysqlConfig) formatDSN() (string, error) {
	dsn, err := mysql.ParseDSN(c.DSN)
	if err != nil {
		return "", err
	}

	// adjust timeout of DSN
	if dsn.Timeout <= 0 {
		dsn.Timeout = c.DialTimeout * time.Millisecond
	}
	if dsn.ReadTimeout <= 0 {
		dsn.ReadTimeout = c.ReadTimeout * time.Millisecond
	}
	if dsn.WriteTimeout <= 0 {
		dsn.WriteTimeout = c.WriteTimeout * time.Millisecond
	}
	return dsn.FormatDSN(), nil
}

func (c *MysqlConfig) getDialector() (gorm.Dialector, error) {
	dsn, err := c.formatDSN()
	if err != nil {
		return nil, err
	}
	return gormysql.Open(dsn), nil
}

func Dial(config *MysqlConfig) (*gorm.DB, error) {
	config.FillWithDefaults()
	d, err := config.getDialector()
	if err != nil {
		return nil, err
	}
	db, err := gorm.Open(d, &gorm.Config{
		Logger: logger.Discard,
	})
	if err != nil {
		return nil, err
	}
	connPool, err := db.DB()
	if err != nil {
		return nil, err
	}
	if config.MaxOpenConns > 0 {
		connPool.SetMaxOpenConns(config.MaxOpenConns)
	}
	if config.MaxIdleConns > 0 {
		connPool.SetMaxIdleConns(config.MaxIdleConns)
	}
	if config.MaxLifeConns > 0 {
		connPool.SetConnMaxLifetime(time.Duration(config.MaxLifeConns) * time.Second)
	}
	err = connPool.Ping()
	if err != nil {
		return nil, err
	}
	return db, nil
}

type Record struct {
	MessageGroupID string  `json:"message_group_id"`
	RowData        RowData `json:"row_data"`
}

type RowData map[string]interface{}

type Table struct {
	MessageGroupID string   `yaml:"message_group_id"`
	Name           string   `yaml:"name"`
	Column         []string `yaml:"column"`
}

func (tbl *Table) GetSelectColumns() string {
	return strings.Join(tbl.Column, ",")
}

func (tbl *Table) GetSQL(offset int, size int) string {
	if offset == 0 {
		return fmt.Sprintf("SELECT %s FROM %s LIMIT %v;",
			tbl.GetSelectColumns(), tbl.Name, size)
	}
	return fmt.Sprintf("SELECT %s FROM %s LIMIT %v OFFSET %v;",
		tbl.GetSelectColumns(), tbl.Name, size, offset)
}

func (tbl *Table) Validate() error {
	if len(tbl.MessageGroupID) == 0 {
		return errors.New("table message_group_id is empty")
	}
	if len(tbl.Name) == 0 {
		return errors.New("table name is empty")
	}
	if len(tbl.Column) == 0 {
		return errors.New("table column is empty")
	}
	return nil
}

type Snapshoter struct {
	db   *gorm.DB
	sink chan Record

	target    []Table
	batchSize int
}

func NewSnapshoter(db *gorm.DB, target []Table, batchSize int) (*Snapshoter, error) {
	return &Snapshoter{
		db:        db,
		sink:      make(chan Record, 120),
		target:    target,
		batchSize: batchSize,
	}, nil
}

// collect query data from database and send data to sink channel.
func (ss *Snapshoter) collect(ctx context.Context) error {
	defer close(ss.sink)
	g, ctx := errgroup.WithContext(ctx)
	for i := range ss.target {
		tbl := ss.target[i]
		g.Go(func() error {
			return ss.drainTable(ctx, tbl)
		})
	}
	return g.Wait()
}

func (ss *Snapshoter) drainTable(ctx context.Context, tbl Table) error {
	offset, cnt := 0, 0
	for cnt < ss.batchSize {
		if err := ctx.Err(); err != nil {
			return err
		}
		sql := tbl.GetSQL(offset, ss.batchSize)
		rows, err := ss.db.Raw(sql).Rows()
		if err != nil {
			log.Printf("fetch data error: %s, SQL: %s\n", err, sql)
			// if table not exist, just return nil
			if mErr, ok := err.(*mysql.MySQLError); ok && mErr.Number == 1146 {
				return nil
			}
			return err
		}
		for rows.Next() {
			data := make(map[string]interface{})
			if err := ss.db.ScanRows(rows, data); err != nil {
				log.Println("scan row error", err)
				return err
			}
			cnt++
			ss.sink <- Record{
				MessageGroupID: tbl.MessageGroupID,
				RowData:        data,
			}
		}
		if cnt == 0 {
			break
		}
		// update offset
		offset, cnt = offset+cnt, 0
	}
	return nil
}

func (ss *Snapshoter) getSink() <-chan Record {
	return ss.sink
}