// 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",
}

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

// 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
  }

  tx, err := me.Tx(ctx, addCpeDictionaryQueryIds)
  if err != nil {
    return err
  }

  // 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
      }
    }
  }

  // commit changes, return result
  return tx.Commit()
}

// 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
  tx, err := me.Tx(ctx, addCpeMatchesQueryIds)
  if err != nil {
    return err
  }

  // 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
      }
    }
  }

  // commit changes, return result
  return tx.Commit()
}

// 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
}