// parses and executes migrations
package migrations

import (
	"code.justin.tv/d8a/iceman/lib/queries"
	"database/sql"
	"errors"
	"fmt"
	yaml "gopkg.in/yaml.v2"
	"strconv"
	"strings"
	"time"
)

type MigrationTarget interface {
	Exec(string, ...interface{}) (sql.Result, error)
}

type Migration struct {
	Up   *MigrationRun
	Down *MigrationRun
}

type MigrationOperation struct {
	Queries []string `yaml:"queries"`
	Timeout string
	Txn     bool

	timeoutValue int
}

type MigrationRun struct {
	Operations []*MigrationOperation
	Timeout    int
}

// helper function to parse yaml
func (m *Migration) parse(data []byte) error {
	return yaml.Unmarshal(data, m)
}

// extract migration from script file
func extractMigration(mr *MigrationRecord) (*Migration, error) {
	data := mr.Content

	var migration Migration
	if err := migration.parse(data); err != nil {
		return &migration, err
	}

	return &migration, nil
}

// ApplyMigration applies (if up) or rolls back (if down) the
// specified migration on the given database
// TODO - A migration record should be distinct from the file which generates it so we can create migration records from other sources.
func ApplyMigration(db *sql.DB, driverQueries queries.DriverQueries, mr *MigrationRecord, direction bool) error {
	migration, err := extractMigration(mr)
	if err != nil {
		return errors.New(fmt.Sprintf("Error processing file \"%v\": %v", mr.Filename, err))
	}

	var run *MigrationRun
	if direction {
		run = migration.Up
	} else {
		run = migration.Down
	}

	//When not using a transaction object, golang could potentially grab a separate
	//transaction for each individual query.  So, it's very bad to do BEGIN or COMMIT manually.
	//Let's make sure nobody's trying to do that.
	err = validateOperations(run)
	if err != nil {
		return errors.New(fmt.Sprintf("Error processing file \"%v\": %v", mr.Filename, err))
	}

	anySucceeded := false
	for _, operation := range run.Operations {
		var target MigrationTarget = db
		var txn *sql.Tx
		opTimeoutValue := run.Timeout

		//Because each query could be made on a separate connection, we
		//should only set timeouts on a transaction when we can use the fruits of our labor
		var txTimeout string
		if operation.Txn {
			txn, err = db.Begin()
			if err != nil {
				return err
			}
			target = txn

			txTimeout = operation.Timeout
			if strings.TrimSpace(txTimeout) != "" {
				opTimeoutValue = operation.timeoutValue
			}

			if opTimeoutValue > 0 {
				_, err := target.Exec(driverQueries.SetTimeout(opTimeoutValue))
				if err != nil {
					return err
				}
			}
		}

		//Execute the queries of this operation
		err = applyToTarget(target, driverQueries, operation.Queries, opTimeoutValue)
		if err != nil {
			return processError(mr, txn, err, operation.Txn, anySucceeded)
		}

		//Commit the transaction if we have one
		if operation.Txn {
			err = txn.Commit()
			if err != nil {
				if anySucceeded {
					fmt.Println("Warning!!! Some transactions committed before the migration returned an error. Check your database.")
				}
				// TODO - This message is repeated in > 1 spot
				return errors.New("Migration failed for " + mr.Name + ", filename:\"" + mr.Filename + "\" (" + err.Error() + ")")
			}
		}
		anySucceeded = true
	}

	//Commit the migration update as a tx
	txn, err := db.Begin()
	if err != nil {
		return err
	}
	if err = updateMigrationTable(txn, driverQueries, direction, mr); err != nil {
		return processError(mr, txn, err, true, anySucceeded)
	}
	return txn.Commit()
}

//Do all the funny error handling we have to do because this logic has gotten a bit complex
// - If we're not in a transaction, warn that some data might have been committed
// - If we're in a transaction but some operations have previously succeeded, warn about that
// - If we're in a transaction, roll it back
// - Return a properly formatted error
func processError(mr *MigrationRecord, txn *sql.Tx, err error, isTxn bool, anySucceeded bool) error {
	if err != nil {
		if !isTxn {
			fmt.Println("Warning!!! A migration failed outside a transaction and some statements may have been committed.  Check your database.")
		} else if anySucceeded {
			fmt.Println("Warning!!! Some transactions committed before the migration returned an error. Check your database.")
		}

		if isTxn {
			e := txn.Rollback()
			if e != nil {
				fmt.Println(e)
			}
		}
		return errors.New("Migration failed for " + mr.Name + ", filename:\"" + mr.Filename + "\" (" + err.Error() + ")")
	}
	return nil
}

//Verify that the migration run doesn't contain any manual tx shenanigans, since database/sql can't handle it
//and ensure there is at least 1 operation to apply
func validateOperations(run *MigrationRun) error {

	if len(run.Operations) == 0 {
		return errors.New("There are no operations to apply.")
	}

	for opIndex, op := range run.Operations {
		for queryIndex, query := range op.Queries {
			trimmedQuery := strings.ToUpper(strings.TrimSpace(query))
			if strings.HasPrefix(trimmedQuery, "BEGIN") || strings.HasPrefix(trimmedQuery, "COMMIT") || strings.HasPrefix(trimmedQuery, "ROLLBACK") {
				return fmt.Errorf("Problem with Operation %d, Query %d: Golang does not support 'loose' transactions.  Add your transaction as a separate operation.", opIndex+1, queryIndex+1)
			}
		}
		if strings.TrimSpace(op.Timeout) != "" {
			timeout, err := strconv.Atoi(strings.TrimSpace(op.Timeout))
			if err != nil {
				return fmt.Errorf("Problem with Operation %d: Timeout must be left blank or be a valid number: %v", opIndex, err)
			}
			op.timeoutValue = timeout
		}
	}

	return nil
}

//Actually apply a migration operation to the DB or tx depending on the sort of operation it is
func applyToTarget(target MigrationTarget, driverQueries queries.DriverQueries, executeQueries []string, timeout int) error {
	var timeoutCh <-chan time.Time
	if timeout > 0 {
		timeoutCh = time.After(time.Duration(timeout) * time.Millisecond)
	} else {
		timeoutCh = make(chan time.Time)
	}

	migrationCh := make(chan error, 1)
	go func() {
		for _, query := range executeQueries {
			if _, err := target.Exec(query); err != nil {
				migrationCh <- err
				return
			}
		}
		migrationCh <- nil
	}()

	select {
	case <-timeoutCh:
		return errors.New("Migration timed out!")
	case err := <-migrationCh:
		return err
	}
}

// Update the version table for the given non-txn migration
func updateMigrationTable(target MigrationTarget, driverQueries queries.DriverQueries, direction bool, mr *MigrationRecord) error {
	var script string

	if direction { // up
		script = driverQueries.InsertMigration()
		utc, err := time.LoadLocation("UTC")
		if err != nil {
			return err
		}
		if _, err := target.Exec(script, mr.Filename, mr.Name, mr.CreatedAt, time.Now().In(utc)); err != nil {
			return err
		}
	} else { // down
		script = driverQueries.DeleteMigration()
		if _, err := target.Exec(script, mr.Filename); err != nil {
			return err
		}
	}
	return nil
}
