package app

import (
	"crypto/tls"
	"crypto/x509"
	"database/sql"
	"fmt"
	"io/ioutil"
	"strings"

	"github.com/pingcap/log"
	"go.uber.org/zap"

	"github.com/go-sql-driver/mysql"
	"github.com/pingcap/errors"
	"github.com/pingcap/tipb/go-tipb"
)

type Store interface {
	WriteCPUTimeRecordBatch([]*tipb.CPUTimeRecord) error
	WriteSQLMetaBatch([]*tipb.SQLMeta) error
	WritePlanMetaBatch([]*tipb.PlanMeta) error
	InitSchema() error
}

var _ Store = &MemStore{}
var _ Store = &TiDBStore{}

// MemStore is for testing purpose
type MemStore struct {
	cpuTimeRecordList []*tipb.CPUTimeRecord
	sqlMetaList       []*tipb.SQLMeta
	planMetaList      []*tipb.PlanMeta
}

func NewMemStore() *MemStore {
	return &MemStore{}
}

func (s *MemStore) WriteCPUTimeRecordBatch(records []*tipb.CPUTimeRecord) error {
	s.cpuTimeRecordList = append(s.cpuTimeRecordList, records...)
	return nil
}

func (s *MemStore) WriteSQLMetaBatch(batch []*tipb.SQLMeta) error {
	s.sqlMetaList = append(s.sqlMetaList, batch...)
	return nil
}

func (s *MemStore) WritePlanMetaBatch(batch []*tipb.PlanMeta) error {
	s.planMetaList = append(s.planMetaList, batch...)
	return nil
}

func (s *MemStore) InitSchema() error {
	return nil
}

// TiDBStore uses TiDB as the storage backend
type TiDBStore struct {
	clusterID    uint64
	instanceID   string
	instanceName string
	instanceComp string
	db           *sql.DB
}

func NewTiDBStore(
	dsn string,
	clusterID uint64,
	instanceID string,
	instanceName string,
	instanceComp string,
	tlsCAPath string,
	tlsCertPath string,
	tlsKeyPath string,
) (*TiDBStore, error) {
	if tlsCAPath != "" || tlsCertPath != "" || tlsKeyPath != "" {
		rootCertPool := x509.NewCertPool()
		pem, err := ioutil.ReadFile(tlsCAPath)
		if err != nil {
			log.Error("failed to read CA cert file", zap.Error(err))
			return nil, err
		}
		if ok := rootCertPool.AppendCertsFromPEM(pem); !ok {
			log.Error("failed to append PEM", zap.Error(err))
			return nil, err
		}
		clientCert := make([]tls.Certificate, 0, 1)
		certs, err := tls.LoadX509KeyPair(tlsCertPath, tlsKeyPath)
		if err != nil {
			log.Error("failed to load x509 key pairs", zap.Error(err))
			return nil, err
		}
		clientCert = append(clientCert, certs)
		if err = mysql.RegisterTLSConfig("topsql", &tls.Config{
			RootCAs:      rootCertPool,
			Certificates: clientCert,
			MinVersion:   tls.VersionTLS12,
		}); err != nil {
			log.Error("failed to register TLS config", zap.Error(err))
			return nil, err
		}
		config, err := mysql.ParseDSN(dsn)
		if err != nil {
			log.Error("failed to parse DSN", zap.Error(err))
			return nil, err
		}
		config.TLSConfig = "topsql"
		dsn = config.FormatDSN()
	}
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		log.Error("failed to open db", zap.Error(err))
		return nil, err
	}
	if err := db.Ping(); err != nil {
		log.Error("failed to ping db", zap.Error(err))
		return nil, err
	}
	log.Info("connecting to TiDB successfully")
	return &TiDBStore{
		clusterID:    clusterID,
		db:           db,
		instanceID:   instanceID,
		instanceName: instanceName,
		instanceComp: instanceComp,
	}, nil
}

func (s *TiDBStore) WriteCPUTimeRecordBatch(batch []*tipb.CPUTimeRecord) error {
	if len(batch) == 0 {
		return nil
	}
	metricsTable := fmt.Sprintf("tidb%d_metrics", s.clusterID)
	values := make([]interface{}, 0, len(batch)*5)
	// iterate through records batch
	for _, record := range batch {
		for j, ts := range record.RecordListTimestampSec {
			cpuTimeMs := record.RecordListCpuTimeMs[j]
			values = append(values, string(record.SqlDigest), string(record.PlanDigest), ts, cpuTimeMs, s.instanceID)
		}
	}
	if len(values) == 0 {
		log.Info("WriteCPUTimeRecordBatch: no records to insert")
		return nil
	}
	log.Info("WriteCPUTimeRecordBatch: inserting records", zap.Int("count", len(values)))
	if len(values)%5 != 0 {
		log.Error("len(values) should be multiple of 5", zap.Int("length", len(values)))
		return errors.New("len(values) should be multiple of 5")
	}
	sqlStrBuf := strings.Builder{}
	sqlStrPrefix := fmt.Sprintf("INSERT INTO %s (sql_digest, plan_digest, timestamp, cpu_time_ms, instance_id) VALUES ", metricsTable)
	sqlStrBuf.Grow(len(sqlStrPrefix) + len("(?, ?, ?, ?, ?),")*len(values)/5)
	sqlStrBuf.WriteString(sqlStrPrefix)
	for i := 0; i < len(values); i += 5 {
		if i > 0 {
			sqlStrBuf.WriteString(",")
		}
		sqlStrBuf.WriteString("(?, ?, ?, ?, ?)")
	}
	sqlStr := sqlStrBuf.String()
	if _, err := s.db.Exec(sqlStr, values...); err != nil {
		log.Error("failed to execute sql", zap.String("sql", sqlStr), zap.Int("values length", len(values)))
		return err
	}
	return nil
}

func (s *TiDBStore) WriteSQLMetaBatch(batch []*tipb.SQLMeta) error {
	sqlMetaTable := fmt.Sprintf("tidb%d_sql_meta", s.clusterID)
	var sqlBuf strings.Builder
	sqlBuf.WriteString(fmt.Sprintf("INSERT IGNORE INTO %s (sql_digest, normalized_sql) VALUES ", sqlMetaTable))
	values := make([]interface{}, 0, len(batch))
	for i, meta := range batch {
		if i > 0 {
			sqlBuf.WriteString(", ")
		}
		sqlBuf.WriteString("(?, ?)")
		values = append(values, meta.SqlDigest, meta.NormalizedSql)
	}
	sqlStr := sqlBuf.String()
	if _, err := s.db.Exec(sqlStr, values...); err != nil {
		return err
	}
	return nil
}

func (s *TiDBStore) WritePlanMetaBatch(batch []*tipb.PlanMeta) error {
	sqlMetaTable := fmt.Sprintf("tidb%d_plan_meta", s.clusterID)
	var sqlBuf strings.Builder
	sqlBuf.WriteString(fmt.Sprintf("INSERT IGNORE INTO %s (plan_digest, normalized_plan) VALUES ", sqlMetaTable))
	values := make([]interface{}, 0, len(batch))
	for i, meta := range batch {
		if i > 0 {
			sqlBuf.WriteString(", ")
		}
		sqlBuf.WriteString("(?, ?)")
		values = append(values, meta.PlanDigest, meta.NormalizedPlan)
	}
	sqlStr := sqlBuf.String()
	if _, err := s.db.Exec(sqlStr, values...); err != nil {
		return err
	}
	return nil
}

func (s *TiDBStore) InitSchema() error {
	// TODO(dragonly): partition table maintenance
	metricsTable := fmt.Sprintf("tidb%d_metrics", s.clusterID)
	instanceTable := fmt.Sprintf("tidb%d_instance", s.clusterID)
	sqlMetaTable := fmt.Sprintf("tidb%d_sql_meta", s.clusterID)
	planMetaTable := fmt.Sprintf("tidb%d_plan_meta", s.clusterID)
	if _, err := s.db.Exec(fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s (
		sql_digest VARBINARY(32) NOT NULL,
		plan_digest VARBINARY(32),
		timestamp INTEGER NOT NULL,
		cpu_time_ms INTEGER NOT NULL,
		instance_id VARCHAR(32) NOT NULL
		);`, metricsTable)); err != nil {
		log.Error("failed to create table", zap.Error(err))
		return err
	}
	if _, err := s.db.Exec(fmt.Sprintf("ALTER TABLE %s SET TIFLASH REPLICA 1;", metricsTable)); err != nil {
		log.Error("failed to set tiflash replica", zap.Error(err))
		return err
	}

	if _, err := s.db.Exec(fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s (
		id BIGINT AUTO_RANDOM NOT NULL,
		sql_digest VARBINARY(32) NOT NULL,
		normalized_sql LONGTEXT NOT NULL,
		PRIMARY KEY (id),
		INDEX i_sql_digest (sql_digest),
		UNIQUE (sql_digest)
	);`, sqlMetaTable)); err != nil {
		log.Error("failed to create table", zap.Error(err))
		return err
	}

	if _, err := s.db.Exec(fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s (
		id VARCHAR(32) NOT NULL,
		name VARCHAR(32) NOT NULL,
		component VARCHAR(32) NOT NULL,
		PRIMARY KEY (name)
	);`, instanceTable)); err != nil {
		log.Error("failed to create table", zap.Error(err))
		return err
	}
	if _, err := s.db.Exec(fmt.Sprintf(`INSERT IGNORE INTO %s VALUES (?,?,?);`, instanceTable), s.instanceID, s.instanceName, s.instanceComp); err != nil {
		log.Error("failed to insert into table", zap.String("table", instanceTable), zap.Error(err))
		return err
	}

	if _, err := s.db.Exec(fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s (
		id BIGINT AUTO_RANDOM NOT NULL,
		plan_digest VARBINARY(32) NOT NULL,
		normalized_plan LONGTEXT NOT NULL,
		PRIMARY KEY (id),
		INDEX i_plan_digest (plan_digest),
		UNIQUE (plan_digest)
	);`, planMetaTable)); err != nil {
		log.Error("failed to create table", zap.Error(err))
		return err
	}

	return nil
}