package main

import (
	"encoding/json"
	"flag"
	"log"
	"net/http"
	_ "net/http/pprof"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"
)

var (
	root = flag.String("root", filepath.Join(os.Getenv("HOME"), "Maildir"),
		"root directory to watch")
	poll = flag.Duration("poll", time.Second*time.Duration(10),
		"poll interval between new file check")
	addr = flag.String("addr", ":8080",
		"address to listen for HTTP connections")
)

func folderFromPath(path string) string {
	const sep = string(filepath.Separator)
	parts := strings.Split(path, sep)
	var folder string
	if isMailDir(path) {
		folder = parts[len(parts)-2]
	} else {
		folder = parts[len(parts)-1]
	}
	if folder[0] == '.' {
		return strings.Replace(folder[1:], ".", "/", -1)
	}
	return "Inbox"
}

// isMailDirParent determines if path is a Maildir folder's
// cur/new/tmp parent directory.
func isMailDirParent(path string) bool {
	if isMailDir(path) {
		return false
	}
	fi, err := os.Stat(filepath.Join(path, "new"))
	if err != nil {
		return false
	}
	return fi.IsDir()
}

// isMailDir determins if path is one of Maildir's cur/new/tmp mail folders.
func isMailDir(path string) bool {
	return strings.HasSuffix(path, "new") || strings.HasSuffix(path, "cur") || strings.HasSuffix(path, "tmp")
}

// isNewEmail parses path looking for Maildir filename patterns for the flags
// indicating if a message has been seen.
func isNewEmail(path string) bool {
	// More docs on Maildir flags at http://cr.yp.to/proto/maildir.html
	//
	// Seen:
	// 1340520791.M843511P28536.sagan.sf.xinu.tv,S=1467,W=1503:2,RS
	// Unseen:
	// 1377995543.50996_2.sagan.sf.xinu.tv:2,Rb
	seen := false
	i := strings.LastIndex(path, ":")
	if i != -1 {
		info := path[i:]
		i = strings.Index(info, ",")
		if i == -1 {
			return true
		}

		flags := info[i:]
		seen = strings.Contains(flags, "S")
	}
	return !seen
}

type mailCounter struct {
	lastrun    time.Time
	counts     map[string]int
	needUpdate []string
	root       string
}

func newMailCounter() *mailCounter {
	return &mailCounter{
		counts: map[string]int{},
	}
}

func (mc *mailCounter) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Access-Control-Allow-Origin", "*")
	c := map[string]int{}
	for k, v := range mc.counts {
		if v != 0 {
			c[k] = v
		}
	}
	enc := json.NewEncoder(w)
	err := enc.Encode(c)
	if err != nil {
		log.Println("Error printing JSON:", err)
	}
}

func (mc *mailCounter) updateCount(path string) error {
	if !isMailDir(path) {
		return nil
	}

	d, err := os.Open(path)
	if err != nil {
		return err
	}

	fi, err := d.Stat()
	if err != nil {
		return err
	}

	if !fi.IsDir() {
		return nil
	}

	folder := folderFromPath(filepath.Dir(path))
	paths, err := d.Readdirnames(-1)
	for _, path := range paths {
		if isNewEmail(path) {
			mc.counts[folder]++
		}
	}
	return nil
}

func (mc *mailCounter) walk(path string, info os.FileInfo, err error) error {
	if err != nil {
		return err
	}

	if info.IsDir() {
		return nil
	}

	dn := filepath.Dir(path)
	if isMailDir(dn) {
		if isNewEmail(path) {
			folder := folderFromPath(dn)
			mc.counts[folder]++
		}
	}
	return nil
}

func skip(path string) bool {
	switch filepath.Base(path) {
	case "cur", "new", "tmp":
		return true
	}
	return false
}

func (mc *mailCounter) fillNeedUpdate() error {
	mc.needUpdate = mc.needUpdate[:0]
	//	log.Println("BEGIN fillNeedUpdate", mc.needUpdate)
	//	defer func() {
	//		log.Println("END fillNeedUpdate", mc.needUpdate)
	//	}()
	return filepath.Walk(mc.root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if info.IsDir() {
			if info.ModTime().After(mc.lastrun) && isMailDir(path) {
				mc.needUpdate = append(mc.needUpdate, path)
			}
			if skip(path) {
				return filepath.SkipDir
			}
		}
		return nil
	})
}

func (mc *mailCounter) watch(root string, poll time.Duration) error {
	mc.root = root
	for {
		thisrun := time.Now()
		if err := mc.fillNeedUpdate(); err != nil {
			return err
		}
		modCnt := len(mc.needUpdate)
		if modCnt > 0 {
			log.Println(modCnt, "folders need update")
			update := map[string]struct{}{}
			for _, dir := range mc.needUpdate {
				// Store parent directory, which is the IMAP folder
				update[filepath.Dir(dir)] = struct{}{}
			}
			for u := range update {
				folder := folderFromPath(u)
				mc.counts[folder] = 0
			}
			for u := range update {
				if err := mc.updateCount(filepath.Join(u, "cur")); err != nil {
					return err
				}
				if err := mc.updateCount(filepath.Join(u, "new")); err != nil {
					return err
				}
				//if err := filepath.Walk(filepath.Join(u, "cur"), mc.walk); err != nil {
				//	return err
				//}
				//if err := filepath.Walk(filepath.Join(u, "new"), mc.walk); err != nil {
				//	return err
				//}
			}
		}
		mc.lastrun = thisrun
		time.Sleep(poll)
	}
	return nil
}

type countList []count

type count struct {
	f string
	c int
}

func (cl countList) Len() int { return len(cl) }
func (cl countList) Less(i, j int) bool {
	if cl[i].c == cl[j].c {
		return cl[i].f < cl[j].f
	}
	return cl[i].c > cl[j].c
}
func (cl countList) Swap(i, j int) { cl[i], cl[j] = cl[j], cl[i] }

func countListFromMap(counts map[string]int) countList {
	cl := make([]count, len(counts))
	i := 0
	for k, v := range counts {
		cl[i].f = k
		cl[i].c = v
		i++
	}
	return cl
}

// printCounts prints the folders with unread messages in descending order,
// sorting alphabetically for tie-breakers
func printCounts(counts map[string]int) {
	cl := countListFromMap(counts)
	sort.Sort(cl)
	for _, c := range cl {
		if c.c > 0 {
			log.Printf("%s (%d)", c.f, c.c)
		}
	}
}

func main() {
	flag.Parse()

	mc := newMailCounter()
	go func() {
		http.Handle("/unread", mc)
		err := http.ListenAndServe(*addr, nil)
		if err != nil {
			log.Fatal(err)
		}
	}()

	err := mc.watch(*root, *poll)
	if err != nil {
		log.Fatal(err)
	}
}