/*
Description: Tool to save bookmarks.
Auth0r: sml@lacashita.com

Use it only as educational purpose.
*/

package main

import (
    "flag"
    "fmt"
    "database/sql"
	"log"
	"os"
	"github.com/fatih/color"
	_ "github.com/mattn/go-sqlite3"
    "github.com/rodaine/table"
    "strconv"
    "time"
    "github.com/chromedp/chromedp"
    "crypto/md5"
    "encoding/hex"
    "io/ioutil"
    "context"
)

func check(e error) {
    if e != nil {
        panic(e)
    }
}


const help = `

(·)_(o) babil0nia-kli

Basic help
-----------
Add url: bkli add -u https://lacashita.com -d "Cool blog"
Add url with tags: bkli add -u https://lacashita.com -d "Cool blog" -t linux,coding,go
List url by tag: bkli list -t linux
List url by word: bkli list -w myword
List all urls: bkli list
List url by status: bkli list -s Open 
Delete url: bkli delete <id>
Modify tags: bkli mod <id> -t hacking,testing
Modify description : bkli mod <id> -d "My new description"
Download Url: bkli download <id>
View image if downloaded: bkli view <id>


Check more help in README.md file.
`

func main() {
    // Get babil0nia_DB env, if doesnt exists create ~/babil0nia.db as database.
    var filedb string
    var homepath = os.Getenv("HOME")
    if os.Getenv("babil0nia_DB") == "" {
        filedb = homepath+"/babil0nia.db"
    } else {
        filedb = os.Getenv("babil0nia_DB")
    }
    // Create ~/.babil0niafiles folder to save screenshots
    if _, err := os.Stat(homepath+"/.babil0niafiles"); os.IsNotExist(err) {
        err = os.Mkdir(homepath+"/.babil0niafiles", 0755)
        if err != nil {
        log.Fatal(err)
        }
    }
    checkDB(filedb)
    bklimain(len(os.Args),filedb)
}

type Url struct {
id *int
date *string
url *string
description *string
tags *string 
screenshot *string 
}

// Main function.
func bklimain(arguments int,filedb string){
    var u Url
    var existid int

    sqlDB, _ := sql.Open("sqlite3",filedb)
    defer sqlDB.Close()

    if len(os.Args) < 2 {
        fmt.Println(help)
        os.Exit(0)
    }

    switch os.Args[1] {

    case "add":
        addCommand := flag.NewFlagSet("add", flag.ExitOnError)
        u.url = addCommand.String("u", "", "Url")
        u.description = addCommand.String("d", "", "Description")
        u.tags = addCommand.String("t", "", "Tags")
        addCommand.Parse(os.Args[2:])
        var verify int 

        if *u.url == "" || *u.description == "" {
            fmt.Println("Parameter -u or -d should not be empty")
            os.Exit(0)
        } else {
        // Check if the URL to add exists.
        verify = checkDupURL(sqlDB,*u.url)
        if verify == 0 {
        addUrl(sqlDB,*u.url, *u.description, *u.tags)
        } else {
            fmt.Println("URL exists.")
            os.Exit(0)
        }
    }   
        
    case "list":
        listCommand := flag.NewFlagSet("list", flag.ExitOnError)
        u.description = listCommand.String("w", "%", "Word")
        u.tags = listCommand.String("t", "%", "Tag")

        listCommand.Parse(os.Args[2:])
        listUrl(sqlDB,*u.description,*u.tags)

    case "mod":
        if len(os.Args) > 2{
            modCommand := flag.NewFlagSet("mod", flag.ExitOnError)
            u.url = modCommand.String("u", "", "URL")
            u.description = modCommand.String("d", "", "Description")
            u.tags = modCommand.String("t", "", "Tags")
            modCommand.Parse(os.Args[3:])

            theintid, err := strconv.Atoi(os.Args[2])
            if err != nil {
                fmt.Println("Parameter should be the ID.")
                os.Exit(0)
            }
            existid = checkID(sqlDB,theintid)
            if existid != 0 {
                modURL(sqlDB,*u.url,*u.description,*u.tags,theintid)
            } else {
                fmt.Println("ID dont exist.")
            }
        } else {
            fmt.Println("Usage: bkli mod <id>")
        }

    case "download":
        if len(os.Args) > 2{ 
            // Convert param 2 into int.
            theintid, err := strconv.Atoi(os.Args[2])
            if err != nil {
                fmt.Println("Parameter should be the ID.")
                os.Exit(0)
            }
            existid = checkID(sqlDB,theintid)
            if existid != 0 {
                urltodownload := obtainURL(sqlDB,theintid)
                downloadedfile := checkDownload(sqlDB,urltodownload)
                filename := getFilename(urltodownload)
                var homepath = os.Getenv("HOME")
                screenshotfullpath := homepath+"/.babil0niafiles/"+filename+".jpg"
                screenshotname := filename+".jpg"
             

                if downloadedfile == false {
                err2 := archiveURL(urltodownload,screenshotfullpath,screenshotname,sqlDB)
                    if err2 != nil {
                        panic(err2)
                    }
                } else {
                    fmt.Println("URL downloaded in the past.")
                }
            } else {
                fmt.Println("ID dont exist.")
            }
        } else {
            fmt.Println("Usage: bkli download <id>")
        }

    case "delete":
        if len(os.Args) > 2{ 
            // Convert param 2 into int.
            theintid, err := strconv.Atoi(os.Args[2])
            if err != nil {
                fmt.Println("Parameter should be the ID.")
                os.Exit(0)
            }
            existid = checkID(sqlDB,theintid)
            if existid != 0 {
                filetodel := obtainDownload(sqlDB,theintid)
                if filetodel != "" {
                    var homepath = os.Getenv("HOME")
                    filetodelete := homepath+"/.babil0niafiles/"+filetodel
                    deleteFile(filetodelete)
                }
                deleteURL(sqlDB,theintid)
            } else {
                fmt.Println("ID dont exist.")
            }
        } else {
            fmt.Println("Usage: bkli delete <id>")
        }

    default:
        fmt.Println(help)
        os.Exit(0)
    }
}

// Func to add Url.
func addUrl(db *sql.DB, url string, description string, tags string) {
    date := time.Now().Local().String()
	insertSQL := `INSERT INTO babil0nia(url, description, tags, date) VALUES (?, ?, ?, ?)`
	statement, err := db.Prepare(insertSQL) 
	if err != nil {
		log.Fatalln(err.Error())
	}
	_, err = statement.Exec(url, description, tags, date)
	if err != nil {
		log.Fatalln(err.Error())
	}
	color.Green("[+] URL added\n")
}


// Func to modify Url.
func modURL(db *sql.DB, url string, description string, tags string, id int) {
// Mod url
if url != "" {
    modSQL := `UPDATE babil0nia set url = ? where id = ?`
    statement, err := db.Prepare(modSQL) 
    if err != nil {
        log.Fatalln(err.Error())
    }
    _, err = statement.Exec(url,id)
    if err != nil {
        log.Fatalln(err.Error())
    }
}
// Mod description
if description != "" {
	modSQL := `UPDATE babil0nia set description = ? where id = ?`
	statement, err := db.Prepare(modSQL) 
	if err != nil {
		log.Fatalln(err.Error())
	}
	_, err = statement.Exec(description,id)
	if err != nil {
		log.Fatalln(err.Error())
	}
}
// Mod tags
if tags != "" {
	modSQL := `UPDATE babil0nia set tags = ? where id = ?`
	statement, err := db.Prepare(modSQL) 
	if err != nil {
		log.Fatalln(err.Error())
	}
	_, err = statement.Exec(tags,id)
	if err != nil {
		log.Fatalln(err.Error())
	}
}

//If all ok prints:
color.Green("[+] URL modified\n")
}

// Func to show all URLs.
func listUrl(db *sql.DB, description string, tags string) {
    headerFmt := color.New(color.FgMagenta, color.Underline).SprintfFunc()
    columnFmt := color.New(color.FgGreen).SprintfFunc()
    tbl := table.New("ID", "Url", "Description", "Tags","Screenshot")
    tbl.WithHeaderFormatter(headerFmt).WithFirstColumnFormatter(columnFmt)
	row, err := db.Query("SELECT * FROM babil0nia where id like '%' and description like ? and tags like ? ","%"+description+"%","%"+tags+"%")
	if err != nil {
		log.Fatal(err)
	}
	defer row.Close()
    fmt.Printf("\n")
	for row.Next() {
		var id int
		var date string
		var url string
        var description string
        var tags string
        var screenshot string
		row.Scan(&id, &date, &url, &description, &tags, &screenshot)
        tbl.AddRow(id,url,description,tags,screenshot)
	}

tbl.Print()
fmt.Printf("\n")
}

// Func to delete Url.
func deleteURL(db *sql.DB, id int) {
	deleteTaskSQL := `DELETE FROM babil0nia WHERE id = ?`
	statement, err := db.Prepare(deleteTaskSQL) 
	if err != nil {
		log.Fatalln(err.Error())
	}
	_, err = statement.Exec(id)
	if err != nil {
		log.Fatalln(err.Error())
	}
	color.Green("[+] URL %v deleted successfully.\n",id)
}

// Func to check if ID exists
func checkID(db *sql.DB, id int ) int{
    var total int
    row, err := db.Query("SELECT count(*) as total from babil0nia where id = ?",id)
    if err != nil {
        log.Fatal(err)
    }
    defer row.Close()
    for row.Next() {
        var count int
        row.Scan( &count)
        total = count   
    }
    return total
}

// Func to check if Url is duplicated
func checkDupURL(db *sql.DB, url string ) int{
    var total int
    row, err := db.Query("SELECT count(*) as total from babil0nia where url = ?",url)
    if err != nil {
        log.Fatal(err)
    }
    defer row.Close()
    for row.Next() {
        var count int
        row.Scan( &count)
        total = count   
    }
    return total
}

// Func to check if url has been downloaded.
func checkDownload(db *sql.DB, url string ) bool {
    downloadedfile := false
    row, err := db.Query("SELECT screenshot from babil0nia where url = ?",url)
    if err != nil {
        log.Fatal(err)
    }
    defer row.Close()
    for row.Next() {
        var screenshot string
        row.Scan( &screenshot)
  
        if screenshot != "" {
        {
            downloadedfile = true
        }
    }
}
return downloadedfile
}

// Func to check if DB file exists, if not, it creates the file.
func checkDB(filedb string) {
    _, err := os.Stat(filedb)
    //If file htask.db doesnt exist.
    if err != nil {
        color.Green("[+]Creating Database\n")
        file, err := os.Create(filedb)
        if err != nil {
            log.Fatal(err.Error())
        }
        file.Close()
        sqlDB, _ := sql.Open("sqlite3",filedb)
        defer sqlDB.Close()
        createTable(sqlDB)
    }
    }

// Func to create Database file/schema.
func createTable(db *sql.DB) {
	createbabil0niaTableSQL := `CREATE TABLE babil0nia (
		"id" integer NOT NULL PRIMARY KEY AUTOINCREMENT,
		"date" TEXT,
		"url" TEXT,
        "description" TEXT,
        "tags" TEXT,
        "screenshot" TEXT
	  );`

	color.Green("[+]Creating babil0nia table...")
	statement, err := db.Prepare(createbabil0niaTableSQL) 
	if err != nil {
		log.Fatal(err.Error())
	}
	statement.Exec() 
	color.Green("[+]babil0nia DB created")
}

// Func to obtain the url given an ID.
func obtainURL(db *sql.DB, id int) string{
    var result string
    row, err := db.Query("SELECT url from babil0nia where id = ?",id)
    if err != nil {
        log.Fatal(err)
    }
    defer row.Close()
    for row.Next() {
        var url string
        row.Scan( &url)
        result = url   
    }
    return result
}

// Func to obtain the download filename given an ID.
func obtainDownload(db *sql.DB, id int) string{
    var result string
    row, err := db.Query("SELECT screenshot from babil0nia where id = ?",id)
    if err != nil {
        log.Fatal(err)
    }
    defer row.Close()
    for row.Next() {
        var screenshot string
        row.Scan( &screenshot)  
        result = screenshot
    }
    return result
}

// Func to create a filename using the md5 hash of the url name.
func getFilename(text string) string {
    hash := md5.Sum([]byte(text))
    return hex.EncodeToString(hash[:])
 }

// Func to take one screenshot from the given URL and save it to a file. It returns error and the name of the saved (screenshot) file.
func archiveURL(urltoarchive string, screenshotfullpath string,screenshotname string,db *sql.DB) (error){
	var fail error
	// create context
        ctx, cancel := chromedp.NewContext(
                context.Background(),
        //        chromedp.WithDebugf(log.Printf),
        )
        defer cancel()
        // capture screenshot of an element
        var buf []byte
        // capture entire browser viewport, returning png with quality=90
        if err := chromedp.Run(ctx, fullScreenshot(urltoarchive, 100, &buf)); err != nil {
				fail = err
        }
		if fail != nil {
		fmt.Println("Fallo la descarga de la imagen asi que nos e guarda")
		} else {
            updateFilename(db,urltoarchive,screenshotname)
        if err := ioutil.WriteFile(screenshotfullpath, buf, 0o644); err != nil {
			fail = err
		}
	}
		return fail
}

func fullScreenshot(urlstr string, quality int, res *[]byte) chromedp.Tasks {
        return chromedp.Tasks{
                chromedp.Navigate(urlstr),
                chromedp.FullScreenshot(res, quality),
        }
}

// Func to add Task.
func updateFilename(db *sql.DB, url string, filename string, ) {
	insertSQL := `UPDATE babil0nia set screenshot = ? where url = ?`
	statement, err := db.Prepare(insertSQL) 
	if err != nil {
		log.Fatalln(err.Error())
	}
	_, err = statement.Exec(filename, url)
	if err != nil {
		log.Fatalln(err.Error())
	}
	color.Green("[+] Download updated.\n")
}

func deleteFile(filetodelete string){
    e := os.Remove(filetodelete)
    if e != nil {
        fmt.Printf("File %v was not found",filetodelete)
    }
}