package main

import (
	"database/sql"
	_ "github.com/Go-SQL-Driver/MySQL"
	"log"
	"os"
	"time"
)

const (
	TIMEZONE = "UTC"
)

var (
	slogger *log.Logger
)

func init() {
	slogger = log.New(os.Stdout, "[ SQL ]", log.LstdFlags)
}

//Retrieves limit Rawdata entries that are older than tim
//limit <= 0 returns all entries that are older than tim
func fetchRawData(db *sql.DB, cfg *Config, tim time.Time) (rDat []RawData, err error) {

	var prepSel *sql.Stmt
	if cfg.Limit > 0 {
		prepSel, err = db.Prepare("SELECT ip_src,ip_dst,as_src,as_dst,port_src,port_dst,packets,pkt_len_distrib,stamp_inserted FROM " + cfg.RawTable + " WHERE stamp_processed IS NULL AND stamp_inserted < ? LIMIT ?")
	} else {
		prepSel, err = db.Prepare("SELECT ip_src,ip_dst,as_src,as_dst,port_src,port_dst,packets,pkt_len_distrib,stamp_inserted FROM " + cfg.RawTable + " WHERE stamp_processed IS NULL AND stamp_inserted < ?")
	}
	if err != nil {
		slogger.Println("Failed to prepare select")
		return
	}

	var rows *sql.Rows
	if cfg.Limit > 0 {
		rows, err = prepSel.Query(tim, cfg.Limit)
	} else {
		rows, err = prepSel.Query(tim)
	}
	if err != nil {
		slogger.Println("Failed to query prepared selection")
		return
	}
	defer rows.Close()

	tx, err := db.Begin()
	if err != nil {
		slogger.Println("Failed to initialize transaction")
		return
	}

	prepUp, err := tx.Prepare("UPDATE " + cfg.RawTable + " SET stamp_processed = ? where ip_src = ? AND ip_dst = ? AND as_src = ? AND as_dst = ? AND port_src = ? AND port_dst = ? AND packets = ? AND pkt_len_distrib = ? AND stamp_inserted = ?")
	if err != nil {
		slogger.Println("Failed to prepare update")
		return
	}

	loc, err := time.LoadLocation(TIMEZONE)
	if err != nil {
		slogger.Println("Couldn't load timezone")
		return
	}
	for rows.Next() {
		var r RawData
		var tim []byte
		err = rows.Scan(&r.Ip_src, &r.Ip_dst, &r.As_src, &r.As_dst, &r.Port_src, &r.Port_dst, &r.Packets, &r.Pkt_len_distrib, &tim)
		if err != nil {
			slogger.Println("Failed to scan result of query")
			return
		}
		r.time, err = time.ParseInLocation("2006-01-02 15:04:05", string(tim), loc)
		if err != nil {
			slogger.Println("Failed to parse timestamp")
			return
		}

		_, err = prepUp.Exec(time.Now(), r.Ip_src, r.Ip_dst, r.As_src, r.As_dst, r.Port_src, r.Port_dst, r.Packets, r.Pkt_len_distrib, tim)
		if err != nil {
			slogger.Println("Failed to query prepared update")
			tx.Rollback()
			return
		}

		rDat = append(rDat, r)
	}
	tx.Commit()
	return
}

//Removes the stamp_processed from every entry that started being proccesed before tim
func reprocess(db *sql.DB, cfg *Config, tim time.Time) (err error) {
	stmt, err := db.Prepare("UPDATE " + cfg.RawTable + " SET stamp_processed = NULL WHERE stamp_processed < ?")
	if err != nil {
		return
	}
	_, err = stmt.Exec(tim)

	return
}

//Removes all entries in the database that started being processed before tim
func purgeAllProcessed(db *sql.DB, cfg *Config, tim time.Time) (err error) {
	stmt, err := db.Prepare("DELETE FROM " + cfg.RawTable + " WHERE stamp_processed < ? ")
	if err != nil {
		return
	}
	_, err = stmt.Exec(tim)

	return
}

//Removes all Rawdata that is in rDat from the database
func purgeRawData(tx *sql.Tx, cfg *Config, rDat []RawData) (err error) {
	prepStmt, err := tx.Prepare("DELETE FROM " + cfg.RawTable + " WHERE ip_src = ? AND ip_dst = ? AND as_src = ? AND as_dst = ? AND port_src = ? AND port_dst = ? AND packets = ? AND pkt_len_distrib = ? AND stamp_processed IS NOT NULL LIMIT 1")
	if err != nil {
		return
	}

	for _, r := range rDat {
		_, err = prepStmt.Exec(r.Ip_src, r.Ip_dst, r.As_src, r.As_dst, r.Port_src, r.Port_dst, r.Packets, r.Pkt_len_distrib)
		if err != nil {
			return
		}
	}
	return
}

func insertCleanData(tx *sql.Tx, cfg *Config, cd []cleanedData) error {
	prepStmt, err := tx.Prepare("INSERT INTO " + cfg.CleanTable + " (ipb_src, ipb_dst, as_src, as_dst, port_src, port_dst, occurences, volume, time_added) VALUES ( ? , ? , ? , ? , ? , ? , ? , ? , ? ) ON DUPLICATE KEY UPDATE occurences = occurences + ?")
	if err != nil {
		slogger.Println("Failed to prepare statement")
		return err
	}

	for ix := range cd {
		_, err = prepStmt.Exec(cd[ix].ipbSrc, cd[ix].ipbDst, cd[ix].asSrc, cd[ix].asDst, cd[ix].portSrc, cd[ix].portDst, cd[ix].occurences, cd[ix].volume, cd[ix].time, cd[ix].occurences)
		if err != nil {
			slogger.Println("Failed to execute statement")
			return err
		}
	}

	return err
}

func insertCleanDataToDB(cfg *Config, cd []cleanedData) error {
	db, err := sql.Open("mysql", cfg.DBUser+":"+cfg.DBPass+"@"+cfg.DBConn+"/"+cfg.DBName)
	if err != nil {
		slogger.Println("Failed to connect to db")
		return err
	}
	defer db.Close()

	prepStmt, err := db.Prepare("INSERT INTO " + cfg.CleanTable + " (ipb_src, ipb_dst, as_src, as_dst, port_src, port_dst, occurences, volume, time_added) VALUES ( ? , ? , ? , ? , ? , ? , ? , ? , ? ) ON DUPLICATE KEY UPDATE occurences = occurences + ?")
	if err != nil {
		slogger.Println("Failed to prepare statement")
		return err
	}

	for ix := range cd {
		_, err = prepStmt.Exec(cd[ix].ipbSrc, cd[ix].ipbDst, cd[ix].asSrc, cd[ix].asDst, cd[ix].portSrc, cd[ix].portDst, cd[ix].occurences, cd[ix].volume, cd[ix].time, cd[ix].occurences)
		if err != nil {
			slogger.Println("Failed to execute statement")
			return err
		}
	}
	return nil
}

func insertASNIP(db *sql.DB, asn int, ipBlock string) error {
	prepCheck, err := db.Prepare("SELECT asn FROM asnip WHERE ip_block = ?")
	if err != nil {
		return err
	}
	defer prepCheck.Close()

	rows, err := prepCheck.Exec(ipBlock)
	if err != nil {
		return err
	}
	if rows != nil {
		return nil
	}

	prepIns, err := db.Prepare("INSERT INTO asnip VALUES ( ? , ? )")
	if err != nil {
		return err
	}
	defer prepIns.Close()

	_, err = prepIns.Exec(asn, ipBlock)
	if err != nil {
		return err
	}

	return nil
}

func removeASNIP(db *sql.DB, asn int, ipBlock string) error {
	prepStmt, err := db.Prepare("DELETE FROM anip WHERE asn = ? AND ip_block = ?")
	if err != nil {
		return err
	}
	defer prepStmt.Close()

	_, err = prepStmt.Exec(asn, ipBlock)
	if err != nil {
		return err
	}

	return nil
}

// Adds differential privacy to all entries in the
// database that is older than t and haven't had
// differential privacy added to them yet.
func privatizeCleaned(db *sql.DB, t time.Time, cfg *Config) (err error) {
	if cfg.Epsilon <= 0 {
		return
	}
	query, err := db.Prepare("SELECT ipb_src,ipb_dst,as_src,as_dst,port_src,port_dst,volume,time_added,occurences FROM " + cfg.CleanTable + " WHERE time_added < ?")
	if err != nil {
		slogger.Println("Failed to prepare query")
		return
	}

	rows, err := query.Query(t)
	if err != nil {
		slogger.Println("Failed to query for unprivitized rows")
		return
	}
	defer rows.Close()

	update, err := db.Prepare("UPDATE " + cfg.CleanTable + " SET occurences = ? , time_privatized = ? WHERE ipb_src = ? AND ipb_dst = ? AND as_src = ? AND as_dst = ? AND port_src = ? AND port_dst = ? AND volume = ? AND time_added = ? ")
	if err != nil {
		slogger.Println("Failed to prepare update")
		return
	}

	loc, err := time.LoadLocation(TIMEZONE)
	if err != nil {
		slogger.Println("Couldn't load timezone")
		return
	}
	var cd cleanedData
	for rows.Next() {
		var tim []byte
		err = rows.Scan(&cd.ipbSrc, &cd.ipbDst, &cd.asSrc, &cd.asDst, &cd.portSrc, &cd.portDst, &cd.volume, &tim, &cd.occurences)
		if err != nil {
			slogger.Println("Failed to scan row")
			return
		}
		cd.time, err = time.ParseInLocation("2006-01-02 15:04:05", string(tim), loc)
		if err != nil {
			slogger.Println("Failed to parse timestamp")
			return
		}
		// Add differential privacy noise
		cd.occurences = diffpriv(cd.occurences, 1, cfg.Epsilon)

		// Update the entry
		_, err := update.Exec(cd.occurences, time.Now(), cd.ipbSrc, cd.ipbDst, cd.asSrc, cd.asDst, cd.portSrc, cd.portDst, cd.volume, cd.time)
		if err != nil {
			slogger.Println("Failed to update an entry:", err)
		}
	}
	return
}

func availableRows(tx *sql.Tx, cfg *Config, timeLimit time.Time) (numRows int, err error) {
	stmt, err := tx.Prepare("SELECT COUNT(*) FROM " + cfg.RawTable + " WHERE stamp_inserted < ? ")
	if err != nil {
		slogger.Println("Could not prepare statement")
		return
	}
	row := stmt.QueryRow(timeLimit)

	err = row.Scan(&numRows)
	if err != nil {
		slogger.Println("Failed to scan result")
	}
	return
}