package db

import (
    "math"
    "time"
    "database/sql"

    _ "github.com/lib/pq"

    "code.justin.tv/common/golibs/log"
)

// ResultMap contains a single db row in a map, with the column name as the
// key for the return value.
type ResultMap map[string]interface{}

// ResultSet is a slice of ResultMap objects, representing the returned results
// of a db query.  For example, assuming the following query:
//    SELECT x, y FROM points;
// the ResultSet would contain one or more more ResultMap objects with two
// keys, "x" and "y".  The type of the associated value interface{} is the
// underlying type returned by the db driver, i.e., the type one would pass to
// golang's sql.Rows.Scan().  They can be easily iterated through using range:
//    points, _ := pool.QueryWithConnection(conn, "SELECT x, y FROM points;")
//    for p := range points {
//        z := math.Sqrt(math.Pow(p["x"].(float64), 2.0) + math.Pow(p["y"].(float64), 2.0))
//    }
type ResultSet []ResultMap

type connectResult struct {
    db  *sql.DB
    err  error
}

// tryClose is sometimes called in situations where the underlying db driver
// may not close the connection; in such instances, the underlying driver may
// panic, or simply return an error without closing.  We will leak resources in
// such situations.  As such, tryClose should be used when it is impossible to
// fully introspect the underlying connection, or where the pool can not
// continue to wait for resources to be returned.  Better handling of this
// situation is future work.
func tryClose(conn *sql.DB) {
    defer func() {
        if r := recover(); r != nil {
            log.Errln("Recoverd error trying to close DB connection: %r", r)
        }
    }()

    if err := conn.Close(); err != nil {
        log.Errln("Error trying to close DB connection: %s", err)
    }
}

func asyncConnect(db_connect string, ret_c chan connectResult) {
    var err  error
    var db  *sql.DB

    db, err = sql.Open("postgres", db_connect)
    if db == nil || err != nil {
        if db != nil {
            tryClose(db)
        }
        ret_c <-connectResult{nil, ConnectionError{err}}
        return
    }

    db.SetMaxIdleConns(1)

    // try a connection; Open does not actually connect
    var result ResultSet
    if result, err = QueryWithConnection(db, "select 1 as dummy;"); err != nil {
        tryClose(db)
        ret_c <-connectResult{nil, ConnectionError{err}}
        return
    }

    if len(result) != 1 {
        tryClose(db)
        ret_c <-connectResult{nil, ImproperlyInitializedError{"DB did not initialize correctly; no results from test query."}}
        return
    }

    if v, ok := result[0]["dummy"].(int64); !ok || v != 1 {
        tryClose(db)
        ret_c <-connectResult{nil, ImproperlyInitializedError{"DB did not initialize correctly; test query has incorrect data."}}
    }

    ret_c <-connectResult{db, err}
}

// NewConnection creates and returns a new connection to the database specified
// by the db_connect parameter string.  The parameter timeout specifies a
// deadline by which the connection must be returned.  Since NewConnection does
// basic tests on the connection before returning, it can return errors from
// both sql.Open and db.QueryWithConnection.  In addition, the following errors
// can be returned:
//
//  * db.ConnectionError: an underlying error was created while attempting to initialize the connection.
//  * db.ConnectionTimeoutError: could not connect within the given duration.
//  * db.ImproperlyInitializedError: db initialized, but returned unexpected bad data; the *sql.DB object is closed and not returned.
func NewConnection(db_connect string, timeout time.Duration) (*sql.DB, error) {
    print_len := int(math.Min(float64(len(db_connect)), 40.0))
    log.Logln("Opening DB connection to \"%s...\"", db_connect[:print_len])

    timer := time.NewTimer(timeout)
    defer timer.Stop()

    connect_c := make(chan connectResult)
    go asyncConnect(db_connect, connect_c)

    select {
    case res := <-connect_c:
        log.Debugln("Returning DB connection from NewConnection.")
        return res.db, res.err

    case <-timer.C:
        log.Errln("NewConnection timed out; not returning new DB connection.")
        go func() {
            res := <-connect_c
            if res.db != nil {
                tryClose(res.db)
            }
            err_s := "None"
            if res.err != nil {
                err_s = res.err.Error()
            }
            log.Errln("Ate late returning DB connect. Error: %s", err_s)
        }()

        return nil, ConnectionTimeoutError{db_connect, timeout}
    }
}

// QueryWithConnection runs the given SQL query on the given *sql.DB object
// with the parameters listed in params.  The types of the parameters must
// match the type expected by the underly SQL driver.  Returns the results of
// the query in a ResultSet and/or an error:
//  * db.QueryFailedError: There was an error in the underlying layers.
//  * db.IntegrityViolationError: Results may have been returned, but there was a SQL integrity violation during the query.
func QueryWithConnection(conn *sql.DB, query string, params ...interface{}) (ResultSet, error) {
    log.Debug("SQL Query (%d params): %s\n", len(params), query)

    // Need to add timeout to this; future work for now.
    rows, err := conn.Query(query, params...)
    if err != nil {
        return nil, QueryFailedError{Query: query, UnderlyingError: err}
    }

    defer rows.Close()

    log.Debugln("SQL Query returned; getting columns.")

    cols, err := rows.Columns()
    if err != nil {
        return nil, QueryFailedError{Query: query, UnderlyingError: err}
    }

    num_cols := len(cols)
    data := make([]interface{}, num_cols)
    for idx := range data  {
        data[idx] = new(interface{})
    }

    log.Debugln("SQL Query scanning.")

    res := make(ResultSet, 0)
    row_idx := 0
    for rows.Next() {
        if err = rows.Scan(data...); err != nil {
            return nil, QueryFailedError{Query: query, UnderlyingError: err}
        }

        val_map := make(ResultMap)
        for i, item := range data  {
            val_map[cols[i]] = *item.(*interface{})
        }
        res = append(res, val_map)

        row_idx += 1
    }

    if err := rows.Err(); err != nil {
        if len(res) > 0 {
            return res, returnQueryError(err, query)
        } else {
            return nil, returnQueryError(err, query)
        }
    }

    log.Debugln("SQL Query done.")

    return res, nil
}

