// database storage
package dbstore

import (
  "context"
  db_sql "database/sql"
  "fmt"
  _ "github.com/mattn/go-sqlite3"
  "github.com/pablotron/cvez/cpedict"
  "github.com/pablotron/cvez/cpematch"
)

// sqlite3 backing store
type DbStore struct {
  db *db_sql.DB
}

// Open database.
//
// This function is called by Open().  It is a separate package-private
// function to make Open() easier to test.
func openFull(dbType, path string) (DbStore, error) {
  var r DbStore

  // init db
  if db, err := db_sql.Open(dbType, path); err != nil {
    return r, err
  } else {
    // save handle
    r.db = db
    return r, nil
  }
}

// Open database
func Open(path string) (DbStore, error) {
  return openFull("sqlite3", path)
}

// initialized database version
const initDbVersion = 314159

func (me DbStore) isInitialized(ctx context.Context) (bool, error) {
  sql := "PRAGMA user_version;"
  // get version
  var version int32
  if err := me.db.QueryRowContext(ctx, sql).Scan(&version); err != nil {
    return false, err
  }

  // return result
  return (version == initDbVersion), nil
}

// enable foreign keys
func (me DbStore) enableForeignKeys(ctx context.Context) error {
  _, err := me.db.ExecContext(ctx, "PRAGMA foreign_keys = true;")
  return err
}

// initialize database
func (me DbStore) Init(ctx context.Context) error {
  if inited, err := me.isInitialized(ctx); err != nil {
    return err
  } else if inited {
    // already initialized, enable foreign keys
    return me.enableForeignKeys(ctx)
  }

  // read init query
  if sql, err := getQuery("init"); err != nil {
    return err
  } else {
    // exec init query, return result
    _, err = me.db.ExecContext(ctx, sql)
    return err
  }
}

var addCpeDictionaryQueryIds = []string {
  "cpe/insert",
  "cpe/insert-title",
  "cpe/insert-ref",
}

// Begin new transaction and create prepared statements.
func (me DbStore) Begin(ctx context.Context, queryIds []string) (Tx, error) {
  return newTx(ctx, me.db, queryIds)
}

// Create a transaction, pass it to callback, then commit the transaction
// if the callback returns success and rollback the transaction if the
// callback returns an error.
func (me DbStore) Tx(ctx context.Context, queryIds []string, fn func(Tx) error) error {
  // create transaction
  tx, err := me.Begin(ctx, queryIds)
  if err != nil {
    return err
  }

  if err := fn(tx); err != nil {
    // rollback
    if rb_err := tx.Rollback(); rb_err != nil {
      return rb_err
    }

    // return error
    return err
  } else {
    // commit transaction
    return tx.Commit()
  }
}

// Execute query and invoke callback with each row of result.
func (me DbStore) Query(
  ctx context.Context,
  queryId string,
  args []interface{},
  fn func(*db_sql.Rows) error,
) error {
  // get query
  sql, err := getQuery(queryId)
  if err != nil {
    return err
  }

  // exec query
  rows, err := me.db.QueryContext(ctx, sql, args...)
  if err != nil {
    return err
  }

  // walk results
  for rows.Next() {
    if err = fn(rows); err != nil {
      return err
    }
  }

  // close rows
  // FIXME: is this correct? i am following the example from the
  // database/sql documentation, but it is messy and it seems
  // counterintuitive to close the row set and then do an additional
  // test for iteration errors...
  if err = rows.Close(); err != nil {
    return err
  }

  // check for iteration errors
  if err = rows.Err(); err != nil {
    return err
  }

  // return success
  return nil
}

// import CPE dictionary
func (me DbStore) AddCpeDictionary(ctx context.Context, dict cpedict.Dictionary) error {
  // lazy-init db
  if err := me.Init(ctx); err != nil {
    return err
  }

  return me.Tx(ctx, addCpeDictionaryQueryIds, func(tx Tx) error {
    // add items
    for _, item := range(dict.Items) {
      // add cpe
      rs, err := tx.Exec(ctx, "cpe/insert", item.CpeUri, item.Cpe23Item.Name)
      if err != nil {
        return err
      }

      // get last row ID
      id, err := rs.LastInsertId()
      if err != nil {
        return err
      }

      // add titles
      for _, title := range(item.Titles) {
        _, err := tx.Exec(ctx, "cpe/insert-title", id, title.Lang, title.Text)
        if err != nil {
          return err
        }
      }

      // add refs
      for _, ref := range(item.References) {
        _, err := tx.Exec(ctx, "cpe/insert-ref", id, ref.Href, ref.Text)
        if err != nil {
          return err
        }
      }
    }

    // return success
    return nil
  })
}

// search CPEs
func (me DbStore) CpeSearch(
  ctx context.Context,
  searchType CpeSearchType,
  s string,
) ([]CpeSearchRow, error) {
  var r []CpeSearchRow

  // lazy-init db
  if err := me.Init(ctx); err != nil {
    return r, err
  }

  // get/exec search query
  err := me.Query(ctx, searchType.String(), []interface{} {
    db_sql.Named("q", s),
  }, func(rows *db_sql.Rows) error {
    if sr, err := unmarshalCpeSearchRow(rows); err != nil {
      // return error
      return err
    } else {
      // append to results
      r = append(r, sr)
      return nil
    }
  })

  // return results
  return r, err
}

// query IDs used by AddCpeMatches()
var addCpeMatchesQueryIds = []string {
  "cpe-match/insert",
  "cpe-match/insert-vulnerable",
  "cpe-match/insert-version-min",
  "cpe-match/insert-version-max",
  "cpe-match/insert-name",
}

// import CPE matches
func (me DbStore) AddCpeMatches(ctx context.Context, matches cpematch.Matches) error {
  // lazy-init db
  if err := me.Init(ctx); err != nil {
    return err
  }

  // begin transaction
  return me.Tx(ctx, addCpeMatchesQueryIds, func(tx Tx) error {
    // add matches
    for _, m := range(matches.Matches) {
      // add cpe
      rs, err := tx.Exec(ctx, "cpe-match/insert", m.Cpe23Uri, m.Cpe22Uri)
      if err != nil {
        return err
      }

      // get last row ID
      id, err := rs.LastInsertId()
      if err != nil {
        return err
      }

      // add vulnerable
      if m.Vulnerable != nil {
        _, err := tx.Exec(ctx, "cpe-match/insert-vulnerable", id, *m.Vulnerable)
        if  err != nil {
          return err
        }
      }

      // add version minimum
      if m.VersionStartIncluding != "" && m.VersionStartExcluding != "" {
        return fmt.Errorf("cannot specify both VersionStartIncluding = \"%s\", VersionEndIncluding \"%s\"", m.VersionStartIncluding, m.VersionStartExcluding)
      } else if m.VersionStartIncluding != "" {
        _, err := tx.Exec(ctx, "cpe-match/insert-version-min", id, true, m.VersionStartIncluding)
        if  err != nil {
          return err
        }
      } else if m.VersionStartExcluding != "" {
        _, err := tx.Exec(ctx, "cpe-match/insert-version-min", id, false, m.VersionStartExcluding)
        if  err != nil {
          return err
        }
      }

      // add version maximum
      if m.VersionEndIncluding != "" && m.VersionEndExcluding != "" {
        return fmt.Errorf("cannot specify both VersionEndIncluding = \"%s\", VersionEndIncluding \"%s\"", m.VersionEndIncluding, m.VersionEndExcluding)
      } else if m.VersionEndIncluding != "" {
        _, err := tx.Exec(ctx, "cpe-match/insert-version-max", id, true, m.VersionEndIncluding)
        if  err != nil {
          return err
        }
      } else if m.VersionEndExcluding != "" {
        _, err := tx.Exec(ctx, "cpe-match/insert-version-max", id, false, m.VersionEndExcluding)
        if  err != nil {
          return err
        }
      }

      // add names
      for _, name := range(m.Names) {
        _, err := tx.Exec(ctx, "cpe-match/insert-name", id, name.Cpe23Uri, name.Cpe22Uri)
        if err != nil {
          return err
        }
      }
    }

    // return success
    return nil
  })
}

// search CPE matches
func (me DbStore) CpeMatchSearch(
  ctx context.Context,
  match string,
) ([]string, error) {
  var r []string

  // lazy-init db
  if err := me.Init(ctx); err != nil {
    return r, err
  }

  // exec search query
  err := me.Query(ctx, "cpe-match/search", []interface{} {
    match,
  }, func(rows *db_sql.Rows) error {
    var s string
    if err := rows.Scan(&s); err != nil {
      // return error
      return err
    } else {
      // append to results
      r = append(r, s)
      return nil
    }
  })

  // return result
  return r, err
}