package sources

import (
	"bytes"
	"database/sql"
	"encoding/csv"
	"encoding/hex"
	"fmt"
	"log"
	"strings"
	"time"

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

	"github.com/tidbcloud/observability/data-collector/pkg/collector/config"
)

const MAX_LINE = 100

type DBSource struct {
	cfg config.DBSource
}

func NewDBSource(cfg config.DBSource) DBSource {

	return DBSource{
		cfg: cfg,
	}
}

var OpenDBfn = func(cfg config.DBSource) (*sql.DB, error) {
	db, err := sql.Open("mysql", cfg.DBUrl)
	if err != nil {
		return nil, err
	}
	return db, nil
}

func (dbs *DBSource) Run(timestamp int64, checkpoint int64) ([]Result, error) {
	// open db
	db, err := OpenDBfn(dbs.cfg)
	if err != nil {
		return nil, err
	}
	defer db.Close()
	results := make([]Result, 0, len(dbs.cfg.Tables))
	// query all tables
	for _, table := range dbs.cfg.Tables {
		var result Result
		buffer := bytes.NewBuffer(nil)
		wcsv := csv.NewWriter(buffer)
		// query count
		rows, err := db.Query(fmt.Sprintf("select count(*) from %s where unix_timestamp(%s)  between %d and %d", table.Name, table.TimeCol, checkpoint, timestamp))
		if err != nil {
			return nil, err
		}

		var count int
		for rows.Next() {
			err = rows.Scan(&count)
			if err != nil {
				return nil, err
			}
		}
		if err = rows.Err(); err != nil {
			return nil, err
		}
		if count == 0 {
			log.Printf("db %s table %s count is 0, skip\n", dbs.cfg.Name, table.Name)
			continue
		}
		log.Printf("db %s table %s count %d\n", dbs.cfg.Name, table.Name, count)
		rows.Close()
		colsName := getColName(table.Cols)
		// query data
		for i := 0; i < count; i += MAX_LINE {
			// sleep 1s for all tables
			time.Sleep(1 * time.Second)
			rows, err := db.Query(fmt.Sprintf("select %s from %s where unix_timestamp(%s) between %d and %d limit %d, %d", strings.Join(colsName, ","), table.Name, table.TimeCol, checkpoint, timestamp, i, MAX_LINE))
			if err != nil {
				return nil, err
			}
			//csv header
			err = wcsv.Write(colsName)
			if err != nil {
				return nil, err
			}
			for rows.Next() {
				values := make([]sql.RawBytes, len(table.Cols))
				cols := make([]interface{}, len(table.Cols))
				for i := range values {
					cols[i] = &values[i]
				}
				// write to csv
				if err := rows.Scan(cols...); err != nil {
					return nil, err
				}
				scols := rawBytes2Strings(values, table.Cols)
				err = wcsv.Write(scols)
				if err != nil {
					return nil, err
				}
			}
			if err = rows.Err(); err != nil {
				return nil, err
			}
			rows.Close()
			wcsv.Flush()

		}
		t := time.Unix(timestamp, 0)
		//make half of hour
		minpart := (t.Minute() / 30)
		filename := fmt.Sprintf("%s-%d-%d_%s.csv", t.Format("20060102-15"), minpart, timestamp, table.Name)
		result.Path = fmt.Sprintf("/%s/%s/%s", dbs.cfg.Name, dbs.cfg.Version, filename)
		result.Body = buffer.Bytes()
		results = append(results, result)
	}
	return results, nil
}

func (dbs *DBSource) GetConfig() config.SourceInterface {
	return dbs.cfg
}

func rawBytes2Strings(values []sql.RawBytes, cols []config.ColType) []string {
	var strcols []string
	for i, v := range cols {
		switch v.Type {
		case "string":
			strcols = append(strcols, string(values[i]))
		case "binary":
			strcols = append(strcols, hex.EncodeToString(values[i]))
		default:
			panic(fmt.Sprintf("unknown type %s", v.Type))
		}
	}
	return strcols
}

func getColName(cols []config.ColType) []string {
	names := make([]string, 0, len(cols))
	for _, col := range cols {
		names = append(names, col.Name)
	}
	return names
}