package db

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"

	"github.com/cheggaaa/pb"
	"github.com/go-redis/redis"
	"github.com/knqyf263/go-cpe/common"
	"github.com/knqyf263/go-cpe/naming"
	c "github.com/kotakanbe/go-cve-dictionary/config"
	log "github.com/kotakanbe/go-cve-dictionary/log"
	"github.com/kotakanbe/go-cve-dictionary/models"
)

/**
# Redis Data Structure

- HASH
  ┌──────────────┬──────────┬──────────┬──────────────────────────────────┐
  │    HASH      │  FIELD   │  VALUE   │             PURPOSE              │
  └──────────────┴──────────┴──────────┴──────────────────────────────────┘
  ┌──────────────┬──────────┬───────────┬─────────────────────────────────┐
  │CVE#${CVEID}  │NVD or JVN│${CVEJSON} │Get CVEJSON by CVEID             │
  ├──────────────┼──────────┼───────────┼─────────────────────────────────┤
  │CVE#C#${CVEID}│NVD or JVN│${CPEJSON} │Get CPEJSON BY CVEID             │
  ├──────────────┼──────────┼───────────┼─────────────────────────────────┤
  │ CVE#Meta     │${URL}    │${METAJSON}│Get FeedMeta BY URL              │
  └──────────────┴──────────┴───────────┴─────────────────────────────────┘

- ZINDE  X
  ┌─────────────────────────┬──────────┬─────────────┬─────────────────────────────────────┐
  │       KEY               │  SCORE   │  MEMBER     │             PURPOSE                 │
  └─────────────────────────┴──────────┴─────────────┴─────────────────────────────────────┘
  ┌─────────────────────────┬──────────┬─────────────┬─────────────────────────────────────┐
  │CVE#${Vendor}::${Product}│    0     │[]${CVEID}   │Get related []CVEID by Vendor,Product│
  └─────────────────────────┴──────────┴─────────────┴─────────────────────────────────────┘

**/

const (
	dialectRedis     = "redis"
	hashKeyPrefix    = "CVE#"
	cpeHashKeyPrefix = "CVE#C#"
)

// RedisDriver is Driver for Redis
type RedisDriver struct {
	name string
	conn *redis.Client
}

// Name return db name
func (r *RedisDriver) Name() string {
	return r.name
}

// NewRedis return Redis driver
func NewRedis(dbType, dbpath string, debugSQL bool) (driver *RedisDriver, locked bool, err error) {
	driver = &RedisDriver{
		name: dbType,
	}
	log.Debugf("Opening DB (%s).", driver.Name())
	if err = driver.OpenDB(dbType, dbpath, debugSQL); err != nil {
		return
	}

	return
}

// OpenDB opens Database
func (r *RedisDriver) OpenDB(dbType, dbPath string, debugSQL bool) (err error) {
	var option *redis.Options
	if option, err = redis.ParseURL(dbPath); err != nil {
		log.Errorf("%s", err)
		return fmt.Errorf("Failed to Parse Redis URL. dbpath: %s, err: %s", dbPath, err)
	}
	r.conn = redis.NewClient(option)
	if err = r.conn.Ping().Err(); err != nil {
		return fmt.Errorf("Failed to open DB. dbtype: %s, dbpath: %s, err: %s", dbType, dbPath, err)
	}
	return nil
}

// CloseDB close Database
func (r *RedisDriver) CloseDB() (err error) {
	if err = r.conn.Close(); err != nil {
		log.Errorf("Failed to close DB. Type: %s. err: %s", r.name, err)
		return
	}
	return
}

// Get Select Cve information from DB.
func (r *RedisDriver) Get(cveID string) (*models.CveDetail, error) {
	var cveResult, cpeResult *redis.StringStringMapCmd
	if cveResult = r.conn.HGetAll(hashKeyPrefix + cveID); cveResult.Err() != nil {
		return nil, cveResult.Err()
	}
	if cpeResult = r.conn.HGetAll(cpeHashKeyPrefix + cveID); cpeResult.Err() != nil {
		return nil, cpeResult.Err()
	}
	return r.unmarshal(cveID, cveResult, cpeResult)
}

// GetMulti Select Cves information from DB.
func (r *RedisDriver) GetMulti(cveIDs []string) (map[string]models.CveDetail, error) {
	cveDetails := map[string]models.CveDetail{}
	pipe := r.conn.Pipeline()
	cveRs, cpeRs := map[string]*redis.StringStringMapCmd{}, map[string]*redis.StringStringMapCmd{}
	for _, cveID := range cveIDs {
		cveRs[cveID] = pipe.HGetAll(hashKeyPrefix + cveID)
		cpeRs[cveID] = pipe.HGetAll(cpeHashKeyPrefix + cveID)
	}
	if _, err := pipe.Exec(); err != nil {
		if err != redis.Nil {
			return nil, fmt.Errorf("Failed to get multi cve json. err : %s", err)
		}
	}

	for cveID, cveResult := range cveRs {
		cpeResult := cpeRs[cveID]
		cveDetail, err := r.unmarshal(cveID, cveResult, cpeResult)
		if err != nil {
			return nil, err
		}
		cveDetails[cveID] = *cveDetail
	}
	return cveDetails, nil
}

func (r *RedisDriver) unmarshal(cveID string, cveResult, cpeResult *redis.StringStringMapCmd) (*models.CveDetail, error) {
	var err error
	jvn := &models.Jvn{}
	if j, ok := cveResult.Val()["Jvn"]; ok {
		if err = json.Unmarshal([]byte(j), jvn); err != nil {
			return nil, err
		}
	}
	if jc, ok := cpeResult.Val()["Jvn"]; ok {
		if err = json.Unmarshal([]byte(jc), &jvn.Cpes); err != nil {
			return nil, err
		}
	}
	if jvn.CveID == "" {
		jvn = nil
	}

	nvdjson := &models.NvdJSON{}
	if j, ok := cveResult.Val()["NvdJSON"]; ok {
		if err = json.Unmarshal([]byte(j), nvdjson); err != nil {
			return nil, err
		}
	}
	if jc, ok := cpeResult.Val()["NvdJSON"]; ok {
		if err = json.Unmarshal([]byte(jc), &nvdjson.Cpes); err != nil {
			return nil, err
		}
	}
	if nvdjson.CveID == "" {
		nvdjson = nil
	}

	nvdxml := &models.NvdXML{}
	if nvdjson == nil {
		if j, ok := cveResult.Val()["Nvd"]; ok {
			if err = json.Unmarshal([]byte(j), nvdxml); err != nil {
				return nil, err
			}
		}
		if jc, ok := cpeResult.Val()["Nvd"]; ok {
			if err = json.Unmarshal([]byte(jc), &nvdxml.Cpes); err != nil {
				return nil, err
			}
		}
		if nvdxml.CveID == "" {
			nvdxml = nil
		}
	} else {
		nvdxml = nil
	}

	return &models.CveDetail{
		CveID:   cveID,
		NvdXML:  nvdxml,
		NvdJSON: nvdjson,
		Jvn:     jvn,
	}, nil
}

// GetByCpeURI Select Cve information from DB.
func (r *RedisDriver) GetByCpeURI(uri string) ([]models.CveDetail, error) {
	specified, err := naming.UnbindURI(uri)
	if err != nil {
		return nil, err
	}
	vendor := fmt.Sprintf("%s", specified.Get(common.AttributeVendor))
	product := fmt.Sprintf("%s", specified.Get(common.AttributeProduct))
	key := fmt.Sprintf("%s%s::%s", hashKeyPrefix, vendor, product)

	var result *redis.StringSliceCmd
	if result = r.conn.ZRange(key, 0, -1); result.Err() != nil {
		return nil, result.Err()
	}

	uniqCveIDs := map[string]bool{}
	for _, v := range result.Val() {
		uniqCveIDs[v] = true
	}
	details := []models.CveDetail{}
	for cveID := range uniqCveIDs {
		d, err := r.Get(cveID)
		if err != nil {
			return nil, err
		}
		match, err := matchCpe(specified, d)
		if err != nil {
			log.Warnf("Failed to compare the version:%s %s %#v",
				err, uri, d)
			// continue matching
			continue
		}
		if match {
			details = append(details, *d)
		}
	}
	return details, nil
}

// InsertJvn insert items fetched from JVN.
func (r *RedisDriver) InsertJvn(cves []models.CveDetail) error {
	log.Infof("Inserting fetched CVEs...")
	var err error
	var refreshedJvns []string
	bar := pb.New(len(cves))
	if c.Conf.Quiet {
		bar.SetWriter(ioutil.Discard)
	} else {
		bar.SetWriter(os.Stderr)
	}
	bar.Start()

	for chunked := range chunkSlice(cves, 10) {
		var pipe redis.Pipeliner
		pipe = r.conn.Pipeline()
		for _, c := range chunked {
			bar.Increment()

			cpes := make([]models.Cpe, len(c.Jvn.Cpes))
			copy(cpes, c.Jvn.Cpes)
			c.Jvn.Cpes = nil

			var jj []byte
			if jj, err = json.Marshal(c.Jvn); err != nil {
				return fmt.Errorf("Failed to marshal json. err: %s", err)
			}
			refreshedJvns = append(refreshedJvns, c.CveID)
			if result := pipe.HSet(hashKeyPrefix+c.CveID, "Jvn", string(jj)); result.Err() != nil {
				return fmt.Errorf("Failed to HSet CVE. err: %s", result.Err())
			}

			for _, cpe := range cpes {
				if result := pipe.ZAdd(
					fmt.Sprintf("%s%s::%s", hashKeyPrefix, cpe.Vendor, cpe.Product),
					redis.Z{Score: 0, Member: c.CveID},
				); result.Err() != nil {
					return fmt.Errorf("Failed to ZAdd cpe. err: %s", result.Err())
				}
			}
			var jc []byte
			if jc, err = json.Marshal(cpes); err != nil {
				return fmt.Errorf("Failed to marshal json. err: %s", err)
			}
			if result := pipe.HSet(cpeHashKeyPrefix+c.CveID, "Jvn", string(jc)); result.Err() != nil {
				return fmt.Errorf("Failed to HSet CPE. err: %s", result.Err())
			}
		}
		if _, err = pipe.Exec(); err != nil {
			return fmt.Errorf("Failed to exec pipeline. err: %s", err)
		}
	}
	bar.Finish()
	log.Infof("Refreshed %d Jvns.", len(refreshedJvns))
	//  log.Debugf("%v", refreshedJvns)
	return nil
}

// CountNvd count nvd table
func (r *RedisDriver) CountNvd() (int, error) {
	var result *redis.StringSliceCmd
	if result = r.conn.Keys(hashKeyPrefix + "CVE*"); result.Err() != nil {
		return 0, result.Err()
	}
	return len(result.Val()), nil
}

// InsertNvdXML inserts CveInformation into DB
func (r *RedisDriver) InsertNvdXML(cves []models.CveDetail) error {
	log.Infof("Inserting CVEs...")
	var err error
	var refreshedNvds []string
	bar := pb.New(len(cves))
	if c.Conf.Quiet {
		bar.SetWriter(ioutil.Discard)
	} else {
		bar.SetWriter(os.Stderr)
	}
	bar.Start()

	for chunked := range chunkSlice(cves, 10) {
		var pipe redis.Pipeliner
		pipe = r.conn.Pipeline()
		for _, c := range chunked {
			bar.Increment()

			cpes := make([]models.Cpe, len(c.NvdXML.Cpes))
			copy(cpes, c.NvdXML.Cpes)
			c.NvdXML.Cpes = nil

			var jn []byte
			if jn, err = json.Marshal(c.NvdXML); err != nil {
				return fmt.Errorf("Failed to marshal json. err: %s", err)
			}
			refreshedNvds = append(refreshedNvds, c.CveID)
			if result := pipe.HSet(hashKeyPrefix+c.CveID, "Nvd", string(jn)); result.Err() != nil {
				return fmt.Errorf("Failed to HSet JVN CVE. err: %s", result.Err())
			}

			for _, cpe := range cpes {
				if result := pipe.ZAdd(
					fmt.Sprintf("%s%s::%s", hashKeyPrefix, cpe.Vendor, cpe.Product),
					redis.Z{Score: 0, Member: c.CveID},
				); result.Err() != nil {
					return fmt.Errorf("Failed to ZAdd cpe. err: %s", result.Err())
				}
			}
			var jc []byte
			if jc, err = json.Marshal(cpes); err != nil {
				return fmt.Errorf("Failed to marshal json. err: %s", err)
			}
			if result := pipe.HSet(cpeHashKeyPrefix+c.CveID, "Nvd", string(jc)); result.Err() != nil {
				return fmt.Errorf("Failed to HSet NVD CPE. err: %s", result.Err())
			}
		}
		if _, err = pipe.Exec(); err != nil {
			return fmt.Errorf("Failed to exec pipeline. err: %s", err)
		}
	}
	bar.Finish()

	log.Infof("Refreshed %d Nvds.", len(refreshedNvds))
	//  log.Debugf("%v", refreshedNvds)
	return nil
}

// InsertNvdJSON Cve information from DB.
func (r *RedisDriver) InsertNvdJSON(cves []models.CveDetail) error {
	log.Infof("Inserting CVEs...")
	var err error
	var refreshedNvds []string
	bar := pb.New(len(cves))
	if c.Conf.Quiet {
		bar.SetWriter(ioutil.Discard)
	} else {
		bar.SetWriter(os.Stderr)
	}
	bar.Start()

	for chunked := range chunkSlice(cves, 10) {
		var pipe redis.Pipeliner
		pipe = r.conn.Pipeline()
		for _, c := range chunked {
			bar.Increment()

			cpes := make([]models.Cpe, len(c.NvdJSON.Cpes))
			copy(cpes, c.NvdJSON.Cpes)
			c.NvdJSON.Cpes = nil
			var jn []byte
			if jn, err = json.Marshal(c.NvdJSON); err != nil {
				return fmt.Errorf("Failed to marshal json. err: %s", err)
			}
			refreshedNvds = append(refreshedNvds, c.CveID)
			if result := pipe.HSet(hashKeyPrefix+c.CveID, "NvdJSON", string(jn)); result.Err() != nil {
				return fmt.Errorf("Failed to HSet CVE. err: %s", result.Err())
			}

			for _, cpe := range cpes {
				if result := pipe.ZAdd(
					fmt.Sprintf("%s%s::%s", hashKeyPrefix, cpe.Vendor, cpe.Product),
					redis.Z{Score: 0, Member: c.CveID},
				); result.Err() != nil {
					return fmt.Errorf("Failed to ZAdd cpe. err: %s", result.Err())
				}
			}
			var jc []byte
			if jc, err = json.Marshal(cpes); err != nil {
				return fmt.Errorf("Failed to marshal json. err: %s", err)
			}
			if result := pipe.HSet(cpeHashKeyPrefix+c.CveID, "NvdJSON", string(jc)); result.Err() != nil {
				return fmt.Errorf("Failed to HSet NVD CPE. err: %s", result.Err())
			}
		}
		if _, err = pipe.Exec(); err != nil {
			return fmt.Errorf("Failed to exec pipeline. err: %s", err)
		}
	}
	bar.Finish()

	log.Infof("Refreshed %d Nvds.", len(refreshedNvds))
	//  log.Debugf("%v", refreshedNvds)
	return nil
}

// GetFetchedFeedMeta selects hash in metafile of the year
func (r *RedisDriver) GetFetchedFeedMeta(url string) (*models.FeedMeta, error) {
	var result *redis.StringStringMapCmd
	if result = r.conn.HGetAll(hashKeyPrefix + "Meta"); result.Err() != nil {
		return nil, result.Err()
	}
	meta := &models.FeedMeta{}
	if s, ok := result.Val()[url]; ok {
		if err := json.Unmarshal([]byte(s), meta); err != nil {
			return nil, err
		}
		return meta, nil
	}
	return meta, nil
}

// UpsertFeedHash selects hash in metafile of the year
func (r *RedisDriver) UpsertFeedHash(m models.FeedMeta) error {
	jn, err := json.Marshal(m)
	if err != nil {
		return fmt.Errorf("Failed to marshal json. err: %s", err)
	}

	var pipe redis.Pipeliner
	pipe = r.conn.Pipeline()
	if result := pipe.HSet(hashKeyPrefix+"Meta", m.URL, jn); result.Err() != nil {
		return fmt.Errorf("Failed to HSet META. err: %s", result.Err())
	}
	if _, err := pipe.Exec(); err != nil {
		return fmt.Errorf("Failed to exec pipeline. err: %s", err)
	}
	return nil
}

// GetFetchedFeedMetas selects a list of FeedMeta
func (r *RedisDriver) GetFetchedFeedMetas() (metas []models.FeedMeta, err error) {
	var result *redis.StringStringMapCmd
	if result = r.conn.HGetAll(hashKeyPrefix + "Meta"); result.Err() != nil {
		return nil, result.Err()
	}
	for _, s := range result.Val() {
		m := models.FeedMeta{}
		if err := json.Unmarshal([]byte(s), &m); err != nil {
			return nil, err
		}
		metas = append(metas, m)
	}
	return
}
