package main

import (
	"bytes"
	"encoding/csv"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"regexp"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
)

type Config struct {
	Zfsnap map[string]string    `json:"zfsnap"`
	Box    map[string]BoxConfig `json:"box"`
	Apps   []AppConfig          `json:apps`
	ssh    map[string]*SSHConfig
}

type Location string

type Snapshot string

type AppConfig struct {
	Name         string              `json:"name"`
	Schedule     []string            `json:"schedule"`
	Sources      []Location          `json:"src"`
	Destinations []Location          `json:"dest"`
	Before       map[string]Location `json:"before"`
	After        map[string]Location `json:"after"`
}

type SSHConfig struct {
	signer   ssh.Signer
	config   *ssh.ClientConfig
	client   *ssh.Client
	logged   bool
	name     string
	zfs      map[string]string
	snapshot []Snapshot
}

type BoxConfig struct {
	Addr string `json:"addr"`
	User string `json:"user"`
	Key  string `json:"key"`
}

var (
	cfgFile   = flag.String("config", "config.json", "config file")
	schedFlag = flag.String("schedule", "", "specific schedule")
	testFlag  = flag.Bool("test", true, "test run")
	debugFlag = flag.Bool("debug", true, "debug")
	cfg       Config
)

func (l Location) Box() string {
	s := strings.Split(string(l), `:`)
	return s[0]
}

func (l Location) Path() string {
	s := strings.Split(string(l), `:`)
	return s[1]
}

func (l Location) Valid() bool {
	s := strings.Split(string(l), `:`)
	return len(s) == 2
}

func (s Snapshot) Path() string {
	s2 := strings.Split(string(s), `@`)
	return s2[0]
}

func (s Snapshot) Name() string {
	s2 := strings.Split(string(s), `@`)
	return s2[1]
}

func (s Snapshot) Append(path string) Snapshot {
	s2 := strings.Split(string(s), `@`)
	return Snapshot(s2[0] + "/" + path + "@" + s2[1])
}

//Load config from file
func (c *Config) Load() error {
	b, err := ioutil.ReadFile(*cfgFile)
	if err != nil {
		if *debugFlag {
			log.Printf("Config.Load : ioutil.ReadFile(%s) : %s", *cfgFile, err)
		}
		return err
	}

	err = json.Unmarshal(b, &c)
	if err != nil {
		if *debugFlag {
			log.Printf("Config.Load : json.Unmarshal : %s", err)
		}
		return err
	}

	cfg.ssh = make(map[string]*SSHConfig)
	for k, v := range c.Box {
		s := &SSHConfig{
			logged: false,
			name:   k,
		}
		cfg.ssh[k] = s
		keyRaw, err := ioutil.ReadFile(v.Key)
		if err != nil {
			if *debugFlag {
				log.Printf("Config.Load : ioutil.ReadFile(%s) : %s", k, err)
			}
			return err
		}

		key, err := ssh.ParseRawPrivateKey(keyRaw)
		if err != nil {
			if *debugFlag {
				log.Printf("Config.Load : ssh.ParseRawPrivateKey(%s) : %s", k, err)
			}
			return err
		}

		s.signer, err = ssh.NewSignerFromKey(key)
		if err != nil {
			if *debugFlag {
				log.Printf("Config.Load : ssh.NewSignerFromKey(%s) : %s", k, err)
			}
			return err
		}

		s.config = &ssh.ClientConfig{
			User: v.User,
			Auth: []ssh.AuthMethod{
				ssh.PublicKeys(s.signer),
			},
			HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		}

		s.client, err = ssh.Dial("tcp", v.Addr, s.config)
		if err != nil {
			if *debugFlag {
				log.Printf("Config.Load : ssh.Dial(%s) : %s", k, err)
			}
			return err
		}

		session, err := s.client.NewSession()
		if err != nil {
			if *debugFlag {
				log.Printf("Config.Load : client.NewSession(%s) : %s", k, err)
			}
			return err
		}

		var b bytes.Buffer
		session.Stdout = &b
		err = session.Run("/usr/bin/uname -a")
		if err != nil {
			if *debugFlag {
				log.Printf("Config.Load : client.NewSession(%s) : %s", k, err)
			}
			return err
		}
		if *debugFlag {
			log.Printf("Config.Load : logged into %s : %s", k, b.String())
		}
		session.Close()
		s.logged = true

	}

	for _, app := range c.Apps {
		for _, src := range app.Sources {
			if !src.Valid() {
				return fmt.Errorf("Source not valid : %s", string(src))
			}
			if _, ok := cfg.ssh[src.Box()]; !ok {
				return fmt.Errorf("No box defined for source : %s", string(src))
			}
		}
		for _, dest := range app.Destinations {
			if !dest.Valid() {
				return fmt.Errorf("Destination not valid : %s", string(dest))
			}
			if _, ok := cfg.ssh[dest.Box()]; !ok {
				return fmt.Errorf("No box defined for destination : %s", string(dest))
			}
		}
		for val, before := range app.Before {
			_, err = regexp.Compile(val)
			if err != nil {
				if *debugFlag {
					log.Printf("Config.Load : invalid regex : %s", val)
				}
				return err
			}
			if !before.Valid() {
				return fmt.Errorf("Before not valid : %s", string(before))
			}
			if _, ok := cfg.ssh[before.Box()]; !ok {
				return fmt.Errorf("No box defined for before : %s", string(before))
			}
		}
		for val, after := range app.After {
			_, err = regexp.Compile(val)
			if err != nil {
				if *debugFlag {
					log.Printf("Config.Load : invalid regex : %s", val)
				}
				return err
			}
			if !after.Valid() {
				return fmt.Errorf("After not valid : %s", string(after))
			}
			if _, ok := cfg.ssh[after.Box()]; !ok {
				return fmt.Errorf("No box defined for after : %s", string(after))
			}
		}
	}

	return nil
}

//Close config
func (c *Config) Close() error {
	return nil
}

func (s *SSHConfig) getLastSnapshot(path string) (Snapshot, error) {
	if *debugFlag {
		log.Printf("SSHConfig.getLastSnapshot : Start %s:%s (%d snapshots)", s.name, path, len(s.snapshot))
	}
	var last Snapshot
	for _, v := range s.snapshot {
		if v.Path() == path {
			last = v
		} else {
			if len(string(last)) > 0 {
				return last, nil
			}
		}
	}
	if len(string(last)) > 0 {
		return last, nil
	}
	return last, fmt.Errorf("no snapshot")
}

func (s *SSHConfig) isLastSnapshot(snapshot Snapshot) bool {
	if *debugFlag {
		log.Printf("SSHConfig.isLastSnapshot : Start %s:%s", s.name, string(snapshot))
	}
	_, err := s.getNextSnapshot(snapshot)
	if err != nil {
		return true
	} else {
		return false
	}
}

func (s *SSHConfig) getFirstSnapshot(path string) (Snapshot, error) {
	if *debugFlag {
		log.Printf("SSHConfig.getFirstSnapshot : Start %s:%s", s.name, path)
	}
	var first Snapshot
	for _, v := range s.snapshot {
		if v.Path() == path {
			first = v
			return first, nil
		}
	}
	return first, fmt.Errorf("no snapshot")
}

func (s *SSHConfig) getNextSnapshot(snapshot Snapshot) (Snapshot, error) {
	if *debugFlag {
		log.Printf("SSHConfig.getNextSnapshot : Start %s:%s", s.name, string(snapshot))
	}
	var next Snapshot
	for id, v := range s.snapshot {
		if v == snapshot {
			if len(s.snapshot) > id+1 {
				next = s.snapshot[id+1]
				if next.Path() == snapshot.Path() {
					return next, nil
				} else {
					return next, fmt.Errorf("no snapshot")
				}
			} else {
				return next, fmt.Errorf("no snapshot")
			}
		}
	}
	return next, fmt.Errorf("no snapshot")
}

func (s *SSHConfig) getSnapshotList() error {
	if *debugFlag {
		log.Printf("SSHConfig.getSnapshotList : Start %s", s.name)
	}
	if !s.logged {
		return fmt.Errorf("Client %s not logged in.", s.name)
	}

	session, err := s.client.NewSession()
	if err != nil {
		if *debugFlag {
			log.Printf("SSHConfig.getSnapshotList : client.NewSession(%s) : %s", s.name, err)
		}
		return err
	}

	var b bytes.Buffer
	session.Stdout = &b
	err = session.Run("/usr/sbin/zfs list -H -t snapshot -o name")
	if err != nil {
		if *debugFlag {
			log.Printf("SSHConfig.getSnapshotList : session.Run(%s) : %s", s.name, err)
		}
		return err
	}

	s.snapshot = make([]Snapshot, 0)

	csvReader := csv.NewReader(&b)
	csvReader.Comma = '\t'
	csvReader.FieldsPerRecord = 1

	csvData, err := csvReader.ReadAll()
	if err != nil {
		if *debugFlag {
			log.Printf("SSHConfig.getSnapshotList : csvReader.ReadAll(%s) : %s", s.name, err)
		}
		return err
	}

	for _, rec := range csvData {
		s.snapshot = append(s.snapshot, Snapshot(rec[0]))
	}

	if *debugFlag {
		log.Printf("SSHConfig.getSnapshotList : %s : read %d zfs snapshots", s.name, len(s.snapshot))
	}

	session.Close()

	return nil
}

func (s *SSHConfig) getZFSList() error {
	if *debugFlag {
		log.Printf("SSHConfig.getZFSList : Start %s", s.name)
	}
	if !s.logged {
		return fmt.Errorf("Client %s not logged in.", s.name)
	}

	session, err := s.client.NewSession()
	if err != nil {
		if *debugFlag {
			log.Printf("SSHConfig.getZFSList : client.NewSession(%s) : %s", s.name, err)
		}
		return err
	}

	var b bytes.Buffer
	session.Stdout = &b
	err = session.Run("/sbin/zfs list -H -o name,mountpoint")
	if err != nil {
		if *debugFlag {
			log.Printf("SSHConfig.getZFSList : session.Run(%s) : %s", s.name, err)
		}
		return err
	}

	s.zfs = make(map[string]string)

	csvReader := csv.NewReader(&b)
	csvReader.Comma = '\t'
	csvReader.FieldsPerRecord = 2

	csvData, err := csvReader.ReadAll()
	if err != nil {
		if *debugFlag {
			log.Printf("SSHConfig.getZFSList : csvReader.ReadAll(%s) : %s", s.name, err)
		}
		return err
	}

	for _, rec := range csvData {
		s.zfs[rec[0]] = rec[1]
	}

	if *debugFlag {
		log.Printf("SSHConfig.getZFSList : %s : read %d zfs file systems", s.name, len(s.zfs))
	}

	session.Close()

	return nil

}

func (s *SSHConfig) isZFS(path string) bool {
	if *debugFlag {
		log.Printf("SSHConfig.isZFS : Start %s:%s", s.name, path)
	}
	if len(s.zfs) == 0 {
		err := s.getZFSList()
		if err != nil {
			if *debugFlag {
				log.Printf("SSHConfig.isZFS : s.getZFSList(%s) : %s", s.name, err)
			}
			return false
		}
	}
	_, ok := s.zfs[path]
	if ok {
		return true
	}
	return false
}

func (s *SSHConfig) exec(cmd string) error {
	if *debugFlag {
		log.Printf("SSHConfig.exec : Start %s on %s", cmd, s.name)
	}

	session, err := s.client.NewSession()
	if err != nil {
		if *debugFlag {
			log.Printf("SSHConfig.exec : client(%s).NewSession(%s) : %s", s.name, cmd, err)
		}
		return err
	}

	err = session.Run(cmd)
	if err != nil {
		if *debugFlag {
			log.Printf("SSHConfig.exec : session(%s).Run(%s) : %s", s.name, cmd, err)
		}
		return err
	}

	session.Close()

	return nil
}

func (s *SSHConfig) createZFS(path string) error {
	if *debugFlag {
		log.Printf("SSHConfig.createZFS : Start %s:%s", s.name, path)
	}
	if len(s.zfs) == 0 {
		err := s.getZFSList()
		if err != nil {
			if *debugFlag {
				log.Printf("SSHConfig.createZFS : s.getZFSList(%s) : %s", s.name, err)
			}
			return err
		}
	}
	p := strings.Split(path, `/`)
	var base string
	for _, d := range p {
		if base == "" {
			base = d
		} else {
			base = base + `/` + d
		}
		if _, ok := s.zfs[base]; !ok {

			if *debugFlag {
				log.Printf("SSHConfig.createZFS : Creating %s:%s", s.name, base)
			}

			err := s.exec("/sbin/zfs create -o mountpoint=none " + base)
			if err != nil {
				if *debugFlag {
					log.Printf("SSHConfig.createZFS : s.exec(%s) : %s", s.name, err)
				}
				return err
			}

			s.zfs[base] = "none"
		}
	}

	return nil

}

func (a AppConfig) RunAppSchedule(schedule string) error {
	if *debugFlag {
		log.Printf("RunAppSchedule : Running %s(%s)", a.Name, schedule)
	}
	for _, src := range a.Sources {
		if !cfg.ssh[src.Box()].isZFS(src.Path()) {
			return fmt.Errorf("No path %s on source", string(src))
		}
		for _, dest := range a.Destinations {
			if !cfg.ssh[dest.Box()].isZFS(dest.Path() + "/" + src.Box() + "/" + src.Path()) {
				err := cfg.ssh[dest.Box()].createZFS(dest.Path() + "/" + src.Box() + "/" + src.Path())
				if err != nil {
					if *debugFlag {
						log.Printf("RunAppSchedule : Error creating %s on %s", dest.Path()+"/"+src.Box()+"/"+src.Path(), dest.Box())
					}
					return err
				}
			}
		}
	}
	for k, v := range a.Before {
		re := regexp.MustCompile(k)
		if re.MatchString(schedule) {
			err := cfg.ssh[v.Box()].exec(v.Path())
			if err != nil {
				if *debugFlag {
					log.Printf("RunAppSchedule : Error executing %s", string(v))
				}
				return err
			}
		}
	}
	var refreshSnapshot map[string]bool
	refreshSnapshot = make(map[string]bool)
	for _, v := range a.Sources {
		if *debugFlag {
			log.Printf("RunAppSchedule : taking %s snapshot for %s", schedule, v.Path())
		}
		err := cfg.ssh[v.Box()].exec("/usr/sbin/zfsnap snapshot -p '" + schedule + "-' -a " + cfg.Zfsnap[schedule] + " " + v.Path())
		if err != nil {
			if *debugFlag {
				log.Printf("RunAppSchedule : Error executing zfsnap on %s", string(v))
			}
			return err
		}
		refreshSnapshot[v.Box()] = true
	}
	for _, v := range a.Destinations {
		refreshSnapshot[v.Box()] = true
	}
	for k, _ := range refreshSnapshot {
		if *debugFlag {
			log.Printf("RunAppSchedule : refreshing snapshots for source %s", k)
		}
		err := cfg.ssh[k].getSnapshotList()
		if err != nil {
			if *debugFlag {
				log.Printf("RunAppSchedule : Error getting snapshots on %s", k)
			}
			return err
		}
	}
	for _, src := range a.Sources {
		for _, dest := range a.Destinations {
			if *debugFlag {
				log.Printf("RunAppSchedule : Sending snapshots from %s to %s", string(src), string(dest))
			}
			dLastSnapshot, err := cfg.ssh[dest.Box()].getLastSnapshot(dest.Path() + "/" + src.Box() + "/" + src.Path())
			if err != nil {
				if *debugFlag {
					log.Printf("RunAppSchedule : No snapshot for %s on %s", string(src), dest.Box())
				}
				sFirstSnapshot, err := cfg.ssh[src.Box()].getFirstSnapshot(src.Path())
				if err != nil {
					if *debugFlag {
						log.Printf("RunAppSchedule : No snapshot for %s", string(src))
					}
					return err
				} else {
					if *debugFlag {
						log.Printf("RunAppSchedule : Initializing snapshot on %s from %s", dest.Box(), string(sFirstSnapshot))
					}
					err = cfg.ssh[dest.Box()].exec("/usr/bin/ssh root@" + src.Box() + " /sbin/zfs send " + string(sFirstSnapshot) + " | /sbin/zfs recv -F " + dest.Path() + "/" + src.Box() + "/" + src.Path())
					if err != nil {
						if *debugFlag {
							log.Printf("RunAppSchedule : Initializing snapshot on %s from %s failed (%s)", dest.Box(), string(sFirstSnapshot), err)
						}
						return err
					}
					var sCurrSnapshot Snapshot
					sNextSnapshot := sFirstSnapshot
					for !cfg.ssh[src.Box()].isLastSnapshot(sNextSnapshot) {
						sCurrSnapshot = sNextSnapshot
						sNextSnapshot, err = cfg.ssh[src.Box()].getNextSnapshot(sNextSnapshot)
						if *debugFlag {
							log.Printf("RunAppSchedule : Sending incrementally %s to %s", string(sNextSnapshot), dest.Box())
						}
						err = cfg.ssh[dest.Box()].exec("/usr/bin/ssh root@" + src.Box() + " /sbin/zfs send -I " + string(sCurrSnapshot) + " " + string(sNextSnapshot) + " | /sbin/zfs recv " + dest.Path() + "/" + src.Box() + "/" + src.Path())
						if err != nil {
							if *debugFlag {
								log.Printf("RunAppSchedule : Sending snapshot on %s from %s failed (%s)", dest.Box(), string(sNextSnapshot), err)
							}
							return err
						}
					}
					if *debugFlag {
						log.Printf("RunAppSchedule : All snapshots sent for %s", string(src))
					}
				}
			} else {
				if *debugFlag {
					log.Printf("RunAppSchedule : Last snapshot on %s is %s", dest.Box(), string(dLastSnapshot))
				}
				var sCurrSnapshot Snapshot
				sNextSnapshot := Snapshot(string(dLastSnapshot)[len(string(dest))+2:])
				for !cfg.ssh[src.Box()].isLastSnapshot(sNextSnapshot) {
					sCurrSnapshot = sNextSnapshot
					sNextSnapshot, err = cfg.ssh[src.Box()].getNextSnapshot(sNextSnapshot)
					if *debugFlag {
						log.Printf("RunAppSchedule : Sending incrementally %s to %s", string(sNextSnapshot), dest.Box())
					}
					err = cfg.ssh[dest.Box()].exec("/usr/bin/ssh root@" + src.Box() + " /sbin/zfs send -I " + string(sCurrSnapshot) + " " + string(sNextSnapshot) + " | /sbin/zfs recv " + dest.Path() + "/" + src.Box() + "/" + src.Path())
					if err != nil {
						if *debugFlag {
							log.Printf("RunAppSchedule : Sending snapshot on %s from %s failed (%s)", dest.Box(), string(sNextSnapshot), err)
						}
						return err
					}
				}
			}
		}
	}
	for k, v := range a.After {
		re := regexp.MustCompile(k)
		if re.MatchString(schedule) {
			err := cfg.ssh[v.Box()].exec(v.Path())
			if err != nil {
				if *debugFlag {
					log.Printf("RunAppSchedule : Error executing %s on %s", v.Path(), v.Box())
				}
				return err
			}
		}
	}
	return nil
}

func main() {
	flag.Parse()

	err := cfg.Load()
	if err != nil {
		log.Printf("Cannot load config (%s)", err)
		os.Exit(1)
	}

	schedule := *schedFlag
	if schedule == "" {
		log.Printf("Main : Finding out schedule.")
		now := time.Now()
		if now.Day() == 1 && int(now.Month()) == 1 {
			schedule = "yearly"
		} else if now.Day() == 1 {
			schedule = "monthly"
		} else if now.Weekday().String() == "Monday" {
			schedule = "weekly"
		} else {
			schedule = "daily"
		}
	}

	err = RunSchedule(schedule)
	if err != nil {
		log.Printf("Cannot run schedule (%s)", err)
		os.Exit(1)
	}
}

//RunSchedule run all backup targets where schedule is registered
func RunSchedule(schedule string) error {
	if *debugFlag {
		log.Printf("RunSchedule : Start %s", schedule)
	}
	if _, ok := cfg.Zfsnap[schedule]; !ok {
		return fmt.Errorf("No retention defined for %s schedule", schedule)
	}

	for _, app := range cfg.Apps {
		for _, schedName := range app.Schedule {
			if schedName == schedule {
				err := app.RunAppSchedule(schedule)
				if err != nil {
					if *debugFlag {
						log.Printf("RunSchedule : Error running %s(%s)", app.Name, schedule)
					}
					return err
				}
			}
		}
	}

	return nil
}