package main

import (
	"database/sql"
	_ "github.com/go-sql-driver/mysql"
	"log"
	"time"
)

const (
	TIMEZONE = "UTC"
)

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

	var prepSel *sql.Stmt
	if 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 acct WHERE 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 acct WHERE stamp_inserted < ?")
	}
	if err != nil {
		log.Println("Failed to prepare select")
		return
	}

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

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

	prepUp, err := tx.Prepare("UPDATE acct 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 {
		log.Println("Failed to prepare update")
		return
	}

	loc, err := time.LoadLocation(TIMEZONE)
	for rows.Next() {
		var r RawData
		var tim []byte
		err = rows.Scan(&r.ipSrc, &r.ipDst, &r.asSrc, &r.asDst, &r.portSrc, &r.portDst, &r.packets, &r.pktLenDist, &tim)
		if err != nil {
			log.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 {
			log.Println("Failed to parse timestamp")
			return
		}

		_, err = prepUp.Exec(time.Now(), r.ipSrc, r.ipDst, r.asSrc, r.asDst, r.portSrc, r.portDst, r.packets, r.pktLenDist, r.time)
		if err != nil {
			log.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, tim time.Time) (err error) {
	stmt, err := db.Prepare("UPDATE acct 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, tim time.Time) (err error) {
	stmt, err := db.Prepare("DELETE FROM acct 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, rDat []RawData) (err error) {
	prepStmt, err := tx.Prepare("DELETE FROM acct 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.ipSrc, r.ipDst, r.asSrc, r.asDst, r.portSrc, r.portDst, r.packets, r.pktLenDist)
		if err != nil {
			return
		}
	}
	return
}

func insertCleanData(tx *sql.Tx, cd *CleanData) error {
	prepStmt, err := tx.Prepare("INSERT INTO clean_data (ipb_src, ipb_dst, as_src, as_dst, port_src, port_dst, occurences, volume, time) VALUES ( ? , ? , ? , ? , ? , ? , ? , ? , ? ) ON DUPLICATE KEY UPDATE occurences = occurences + ?")
	if err != nil {
		log.Println("Failed to prepare statement")
		return err
	}

	_, err = prepStmt.Exec(cd.ipbSrc, cd.ipbDst, cd.asSrc, cd.asDst, cd.portSrc, cd.portDst, cd.occurences, cd.volume, cd.time, cd.occurences)
	if err != nil {
		log.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
}