package mysql

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/jmoiron/sqlx"
	"github.com/tidbcloud/pov-service/internal/common/errors"
	"github.com/tidbcloud/pov-service/internal/model"
	"github.com/tidbcloud/pov-service/internal/repo"
)

type WorkloadSample struct {
	ID         int64 `db:"id"`
	WorkloadID int64 `db:"workload_id"`
	Version    int64 `db:"version"`

	CreatedAt time.Time  `db:"created_at"`
	UpdatedAt *time.Time `db:"updated_at"`

	DBName      string `db:"db_name"`
	TableName   string `db:"table_name"`
	ColumnsJSON string `db:"columns_json"`
	RowsJSON    string `db:"rows_json"`

	CreateStatement sql.NullString `db:"create_statement"`
	TableType       sql.NullString `db:"table_type"`
}

const (
	workloadSampleColumns = `
		id,
		workload_id,
		version,
		created_at,
		updated_at,
		db_name,
		table_name,
		columns_json,
		rows_json,
		create_statement,
		table_type
	`
	workloadSampleSchemaColumns = `
		id,
		workload_id,
		version,
		created_at,
		updated_at,
		db_name,
		table_name,
		columns_json,
		table_type
	`
)

type workloadSampleRepo struct {
	db *sqlx.DB
}

func NewWorkloadSampleRepo(db *sqlx.DB) repo.WorkloadSample {
	return &workloadSampleRepo{
		db: db,
	}
}

func (w *workloadSampleRepo) BeginTx(ctx context.Context) (repo.Tx, error) {
	tx, err := w.db.BeginTxx(ctx, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to begin transaction: %w", err)
	}
	return &mysqlTx{tx: tx}, nil
}

func (w *workloadSampleRepo) LockWorkloadSchemaInTx(
	ctx context.Context, tx repo.Tx, workloadID int64,
) (model.DBSampleResult, error) {
	mysqlTx, ok := tx.(*mysqlTx)
	if !ok {
		return model.DBSampleResult{}, fmt.Errorf("invalid transaction type")
	}

	var dbSample []WorkloadSample
	query := fmt.Sprintf(`SELECT %s
		FROM workload_sample
		WHERE workload_id = ?
		FOR UPDATE NOWAIT
		`, workloadSampleSchemaColumns)
	err := mysqlTx.tx.SelectContext(ctx, &dbSample, query, workloadID)
	if err != nil {
		return model.DBSampleResult{}, err
	}

	dbs, version, err := w.dbToModel(dbSample)
	return model.DBSampleResult{
		Databases: dbs,
		Version:   version,
		DBCount:   len(dbs),
	}, err
}

func (w *workloadSampleRepo) CreateInTx(
	ctx context.Context,
	tx repo.Tx,
	workloadID int64,
	dbs []model.DataBase) (int64, error) {
	// there is no data to be inserted
	if len(dbs) == 0 {
		return 0, nil
	}

	// TODO: The batch insert is good for performance, but may bring high memory requirments
	// Modify to smaller batches is need
	query := `INSERT INTO workload_sample (
    		workload_id,
			version,
			created_at,
			updated_at,
			db_name,
			table_name,
			columns_json,
			rows_json,
			create_statement,
			table_type
	) VALUES %s`
	queryPlaceholder, params, version, err := w.params(workloadID, dbs)
	if err != nil {
		return 0, err
	}
	query = fmt.Sprintf(query, queryPlaceholder)

	mysqlTx, ok := tx.(*mysqlTx)
	if !ok {
		return 0, fmt.Errorf("invalid transaction type")
	}
	_, err = mysqlTx.tx.ExecContext(ctx,
		query,
		params...,
	)
	return version, err
}

func (w *workloadSampleRepo) DeleteOutdatedInTx(
	ctx context.Context,
	tx repo.Tx,
	workloadID int64,
	latestVersion int64) error {
	mysqlTx, ok := tx.(*mysqlTx)
	if !ok {
		return fmt.Errorf("invalid transaction type")
	}

	query := `
		DELETE FROM workload_sample
		WHERE workload_id = ?
		AND version < ?
	`
	_, err := mysqlTx.tx.ExecContext(ctx,
		query,
		workloadID, latestVersion,
	)
	return err
}

func (w *workloadSampleRepo) SelectLatestSchema(
	ctx context.Context,
	workloadID int64) (model.DBSampleResult, error) {
	var dbSample []WorkloadSample
	query := fmt.Sprintf(`SELECT %s
		FROM workload_sample
		WHERE workload_id = ?
		AND version = (
	        SELECT max(version)
			FROM workload_sample
			WHERE workload_id = ?)
		`, workloadSampleSchemaColumns)
	err := w.db.SelectContext(ctx, &dbSample, query, workloadID, workloadID)
	if err != nil {
		return model.DBSampleResult{}, err
	}

	dbs, version, err := w.dbToModel(dbSample)
	return model.DBSampleResult{
		Databases: dbs,
		Version:   version,
		DBCount:   len(dbs),
	}, err
}

func (w *workloadSampleRepo) SelectTable(
	ctx context.Context,
	workloadID int64,
	// versionn int64, TODO
	dbName,
	tableName string,
) (model.Table, error) {
	query := fmt.Sprintf(`
		SELECT %s
		FROM workload_sample
		WHERE workload_id = ?
		AND db_name = ?
		AND table_name = ?
		`, workloadSampleColumns)
	var dbSample WorkloadSample
	err := w.db.GetContext(ctx, &dbSample, query, workloadID, dbName, tableName)
	if err != nil {
		if err == sql.ErrNoRows {
			return model.Table{}, errors.NewNotFound("table not found")
		}
		return model.Table{}, err
	}

	var cols []model.Column
	if dbSample.ColumnsJSON != "" {
		if err := json.Unmarshal([]byte(dbSample.ColumnsJSON), &cols); err != nil {
			return model.Table{}, fmt.Errorf("failed to unmarshal columns: %w", err)
		}
	}

	var rows [][]string
	if dbSample.RowsJSON != "" {
		if err := json.Unmarshal([]byte(dbSample.RowsJSON), &rows); err != nil {
			return model.Table{}, fmt.Errorf("failed to unmarshal rows: %w", err)
		}
	}

	return model.Table{
		Name: dbSample.TableName,
		Data: model.Data{
			Columns: cols,
			Rows:    rows,
		},
		Type:            dbSample.TableType.String,
		CreateStatement: dbSample.CreateStatement.String,
	}, nil
}

func (w *workloadSampleRepo) params(
	workloadID int64,
	dbs []model.DataBase) (query string, params []any, version int64, err error) {
	placeholders := make([]string, 0)
	version = time.Now().UTC().Unix()
	for _, db := range dbs {
		for _, table := range db.Tables {
			/*
				   	query := `INSERT INTO workload_sample (
				       		workload_id,
							version,
				   			created_at,
				   			updated_at,
				   			db_name,
				   			table_name,
				   			columns,
				   			rows,
				   			create_statement,
				   			table_type
				   	) VALUES ...
			*/
			placeholders = append(placeholders, "(?, ?, ?, ?, ?, ?, ?, ?, ?, ?)")

			params = append(params, workloadID)
			params = append(params, version)
			params = append(params, time.Now())
			params = append(params, time.Now())
			params = append(params, db.Name)
			params = append(params, table.Name)

			colBytes, err := json.Marshal(table.Columns)
			if err != nil {
				return "", nil, 0, fmt.Errorf("failed to marshal columns: %w", err)
			}
			params = append(params, string(colBytes))

			rowBytes, err := json.Marshal(table.Rows)
			if err != nil {
				return "", nil, 0, fmt.Errorf("failed to marshal rows: %w", err)
			}
			params = append(params, string(rowBytes))
			params = append(params, table.CreateStatement)
			params = append(params, table.Type)
		}
	}
	query = strings.Join(placeholders, ",")
	return
}

func (w *workloadSampleRepo) dbToModel(
	dbSamples []WorkloadSample) ([]model.DataBase, int64, error) {

	var version int64 = 0
	dbMap := make(map[string]model.DataBase)

	for _, dbSample := range dbSamples {
		db, ok := dbMap[dbSample.DBName]
		if !ok {
			db = model.DataBase{
				Name:   dbSample.DBName,
				Tables: make([]model.Table, 0),
			}
		}
		if version == 0 {
			version = dbSample.Version
		}
		var cols []model.Column
		if dbSample.ColumnsJSON != "" {
			if err := json.Unmarshal([]byte(dbSample.ColumnsJSON), &cols); err != nil {
				return nil, version, fmt.Errorf("failed to unmarshal columns: %w", err)
			}
		}

		var rows [][]string
		if dbSample.RowsJSON != "" {
			if err := json.Unmarshal([]byte(dbSample.RowsJSON), &rows); err != nil {
				return nil, version, fmt.Errorf("failed to unmarshal rows: %w", err)
			}
		}

		db.Tables = append(db.Tables, model.Table{
			Name: dbSample.TableName,
			Data: model.Data{
				Columns: cols,
				Rows:    rows,
			},
			Type:            dbSample.TableType.String,
			CreateStatement: dbSample.CreateStatement.String,
		})

		dbMap[dbSample.DBName] = db
	}
	var dbs []model.DataBase
	for _, db := range dbMap {
		dbs = append(dbs, db)
	}
	return dbs, version, nil
}