package main

import (
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

// snapNames are what the store accepts as a snap's name, optionally with an instance key.
var snapNames = regexp.MustCompile(`^[a-z0-9][a-z0-9-]{0,39}(_[a-z0-9]{1,10})?$`)

func checkSnapName(name string) error {
	if !snapNames.MatchString(name) {
		return fmt.Errorf("%q is not a snap name", name)
	}
	return nil
}

func snapName(args map[string]any) (string, error) {
	name, err := text(args, "name")
	if err != nil {
		return "", err
	}
	return name, checkSnapName(name)
}

var plain = []string{"--unicode=never", "--color=never"}

// Where snapd keeps what it knows. Tests point these elsewhere.
var (
	snapsDir = "/var/lib/snapd/snaps"
	lsmFile  = "/sys/kernel/security/lsm"
	snapRoot = "/snap"
)

// UnitState is one unit's state.
type UnitState struct {
	Unit    string `json:"unit"`
	Enabled string `json:"enabled"`
	Active  string `json:"active"`
}

// StatusAnswer is what snapd_status answers.
type StatusAnswer struct {
	Installed   bool        `json:"installed"`
	Version     string      `json:"version,omitempty"`
	Units       []UnitState `json:"units"`
	AppArmor    bool        `json:"kernel_apparmor"`
	ClassicRoot bool        `json:"classic_root"`
	Findings    []string    `json:"findings"`
}

// Status says whether snapd is here and working.
func Status() (StatusAnswer, error) {
	out := StatusAnswer{Units: []UnitState{}, Findings: []string{}}
	r := run(Cmd{Name: "snap", Args: []string{"version"}})
	if r.Error == "not-found" {
		out.Findings = append(out.Findings, "snapd is not installed on this machine")
		return out, nil
	}
	if r.Status != 0 || r.Error != "" {
		return out, failure(Cmd{Name: "snap", Args: []string{"version"}}, r)
	}
	out.Installed = true
	for _, l := range lines(r.Stdout) {
		if f := strings.Fields(l); len(f) >= 2 && f[0] == "snapd" {
			out.Version = f[1]
		}
	}
	for _, u := range []string{"snapd.socket", "snapd.service", "snapd.apparmor.service", "apparmor.service"} {
		// is-enabled and is-active answer on stdout and exit non-zero for "disabled" and "inactive":
		// a state, not a failure.
		en := run(Cmd{Name: "systemctl", Args: []string{"is-enabled", u}})
		ac := run(Cmd{Name: "systemctl", Args: []string{"is-active", u}})
		if en.Error != "" || ac.Error != "" {
			return out, failure(Cmd{Name: "systemctl", Args: []string{"is-enabled", u}}, en)
		}
		out.Units = append(out.Units, UnitState{Unit: u, Enabled: firstLine(en.Stdout), Active: firstLine(ac.Stdout)})
	}
	if b, err := os.ReadFile(lsmFile); err == nil {
		for _, m := range strings.Split(strings.TrimSpace(string(b)), ",") {
			out.AppArmor = out.AppArmor || m == "apparmor"
		}
	}
	_, err := os.Stat(snapRoot)
	out.ClassicRoot = err == nil
	if out.Units[0].Enabled != "enabled" {
		out.Findings = append(out.Findings, "snapd.socket is not enabled: snapd does not start on demand")
	}
	if !out.AppArmor {
		out.Findings = append(out.Findings, "the kernel does not run AppArmor: strict snaps run without their confinement")
	}
	if out.Units[2].Enabled == "enabled" && !out.AppArmor {
		out.Findings = append(out.Findings, "snapd.apparmor.service is enabled with no AppArmor in the kernel: it loads profiles nothing enforces")
	}
	if !out.ClassicRoot {
		out.Findings = append(out.Findings, "/snap does not exist: classic snaps cannot run")
	}
	return out, nil
}

// Snap is one installed revision.
type Snap struct {
	Name      string   `json:"name"`
	Version   string   `json:"version"`
	Revision  string   `json:"revision"`
	Tracking  string   `json:"tracking"`
	Publisher string   `json:"publisher"`
	Notes     []string `json:"notes"`
	Disabled  bool     `json:"disabled"`
}

// ListAnswer is what snapd_list answers.
type ListAnswer struct {
	Snaps    []Snap `json:"snaps"`
	Active   int    `json:"active"`
	Disabled int    `json:"disabled_revisions"`
}

// columns reads a table snap prints: a header line, then whitespace-separated columns, the last
// taking the rest of the line.
func columns(s string, n int) [][]string {
	out := [][]string{}
	for i, l := range lines(s) {
		if i == 0 {
			continue
		}
		f := strings.Fields(l)
		if len(f) < n {
			continue
		}
		if len(f) > n {
			f = append(f[:n-1], strings.Join(f[n-1:], " "))
		}
		out = append(out, f)
	}
	return out
}

// List answers every installed snap and revision.
func List() (ListAnswer, error) {
	r, err := call(Cmd{Name: "snap", Args: append([]string{"list", "--all"}, plain...)})
	if err != nil {
		return ListAnswer{}, err
	}
	out := ListAnswer{Snaps: []Snap{}}
	for _, f := range columns(r.Stdout, 6) {
		s := Snap{Name: f[0], Version: f[1], Revision: f[2], Tracking: f[3], Publisher: strings.TrimRight(f[4], "*"), Notes: []string{}}
		if f[5] != "-" {
			s.Notes = strings.Split(f[5], ",")
		}
		for _, n := range s.Notes {
			s.Disabled = s.Disabled || n == "disabled"
		}
		if s.Disabled {
			out.Disabled++
		} else {
			out.Active++
		}
		out.Snaps = append(out.Snaps, s)
	}
	return out, nil
}

// InfoAnswer is what snapd_info answers.
type InfoAnswer struct {
	Name     string            `json:"name"`
	Fields   map[string]string `json:"fields"`
	Commands []string          `json:"commands"`
	Channels map[string]string `json:"channels"`
}

// Info reads snap info's YAML-like answer: top-level key: value lines, and the commands and
// channels blocks.
func Info(name string) (InfoAnswer, error) {
	r, err := call(Cmd{Name: "snap", Args: append([]string{"info"}, append(plain, name)...)})
	if err != nil {
		return InfoAnswer{}, err
	}
	out := InfoAnswer{Name: name, Fields: map[string]string{}, Commands: []string{}, Channels: map[string]string{}}
	block := ""
	for _, l := range strings.Split(r.Stdout, "\n") {
		if strings.TrimSpace(l) == "" {
			continue
		}
		if !strings.HasPrefix(l, " ") {
			block = ""
			k, v, found := strings.Cut(l, ":")
			if !found {
				continue
			}
			v = strings.TrimSpace(v)
			switch {
			case v == "" || v == "|":
				block = k
			default:
				out.Fields[k] = v
			}
			continue
		}
		t := strings.TrimSpace(l)
		switch block {
		case "commands":
			out.Commands = append(out.Commands, strings.TrimPrefix(t, "- "))
		case "channels":
			if k, v, found := strings.Cut(t, ":"); found {
				out.Channels[k] = strings.Join(strings.Fields(v), " ")
			}
		case "description":
			out.Fields["description"] = strings.TrimSpace(out.Fields["description"] + " " + t)
		}
	}
	return out, nil
}

// Update is one pending refresh.
type Update struct {
	Name      string `json:"name"`
	Version   string `json:"version"`
	Revision  string `json:"revision"`
	Size      string `json:"size"`
	Publisher string `json:"publisher"`
}

// Updates answers what a refresh would change.
func Updates() (map[string]any, error) {
	r, err := call(Cmd{Name: "snap", Args: append([]string{"refresh", "--list"}, plain...)})
	if err != nil {
		return nil, err
	}
	out := []Update{}
	if !strings.Contains(r.Stdout+r.Stderr, "All snaps up to date") {
		for _, f := range columns(r.Stdout, 6) {
			out = append(out, Update{Name: f[0], Version: f[1], Revision: f[2], Size: f[3], Publisher: strings.TrimRight(f[4], "*")})
		}
	}
	return map[string]any{"updates": out, "count": len(out)}, nil
}

// Revision is one revision's file.
type Revision struct {
	Revision string `json:"revision"`
	Bytes    int64  `json:"bytes"`
	Disabled bool   `json:"disabled"`
}

// SnapUsage is the space one snap's revisions take.
type SnapUsage struct {
	Name      string     `json:"name"`
	Bytes     int64      `json:"bytes"`
	Revisions []Revision `json:"revisions"`
}

// DiskAnswer is what snapd_disk_usage answers.
type DiskAnswer struct {
	Dir         string      `json:"dir"`
	TotalBytes  int64       `json:"total_bytes"`
	Reclaimable int64       `json:"disabled_revisions_bytes"`
	Snaps       []SnapUsage `json:"snaps"`
	Note        string      `json:"note"`
}

// DiskUsage measures the snap files and marks the disabled revisions.
func DiskUsage() (DiskAnswer, error) {
	listed, err := List()
	if err != nil {
		return DiskAnswer{}, err
	}
	disabled := map[string]bool{}
	for _, s := range listed.Snaps {
		if s.Disabled {
			disabled[s.Name+"_"+s.Revision] = true
		}
	}
	files, err := filepath.Glob(filepath.Join(snapsDir, "*.snap"))
	if err != nil {
		return DiskAnswer{}, err
	}
	out := DiskAnswer{Dir: snapsDir, Snaps: []SnapUsage{},
		Note: "A disabled revision is kept by snapd for rollback (refresh.retain); removing one is `snap remove --revision`, which no tool here does."}
	by := map[string]*SnapUsage{}
	for _, f := range files {
		info, err := os.Stat(f)
		if err != nil {
			return DiskAnswer{}, err
		}
		base := strings.TrimSuffix(filepath.Base(f), ".snap")
		i := strings.LastIndex(base, "_")
		if i <= 0 {
			continue
		}
		name, rev := base[:i], base[i+1:]
		if by[name] == nil {
			by[name] = &SnapUsage{Name: name, Revisions: []Revision{}}
		}
		d := disabled[base]
		by[name].Revisions = append(by[name].Revisions, Revision{Revision: rev, Bytes: info.Size(), Disabled: d})
		by[name].Bytes += info.Size()
		out.TotalBytes += info.Size()
		if d {
			out.Reclaimable += info.Size()
		}
	}
	for _, u := range by {
		sort.Slice(u.Revisions, func(i, k int) bool {
			a, _ := strconv.Atoi(u.Revisions[i].Revision)
			b, _ := strconv.Atoi(u.Revisions[k].Revision)
			return a < b
		})
		out.Snaps = append(out.Snaps, *u)
	}
	sort.Slice(out.Snaps, func(i, k int) bool { return out.Snaps[i].Bytes > out.Snaps[k].Bytes })
	return out, nil
}

// Services answers the snaps' services.
func Services() (map[string]any, error) {
	r, err := call(Cmd{Name: "snap", Args: append([]string{"services"}, plain...)})
	if err != nil {
		return nil, err
	}
	out := []map[string]string{}
	if !strings.Contains(r.Stdout+r.Stderr, "no services") {
		for _, f := range columns(r.Stdout, 4) {
			out = append(out, map[string]string{"service": f[0], "startup": f[1], "current": f[2], "notes": f[3]})
		}
	}
	return map[string]any{"services": out}, nil
}

// Change is one of snapd's changes, or one task of a change.
type Change struct {
	ID      string `json:"id,omitempty"`
	Status  string `json:"status"`
	Spawn   string `json:"spawn"`
	Ready   string `json:"ready,omitempty"`
	Summary string `json:"summary"`
}

var changeID = regexp.MustCompile(`^[0-9]+$`)

// Changes answers snapd's recent changes, or one change's tasks.
func Changes(id string) (map[string]any, error) {
	args := append([]string{"changes", "--abs-time"}, plain...)
	n := 5
	if id != "" {
		if !changeID.MatchString(id) {
			return nil, fmt.Errorf("%q is not a change id", id)
		}
		args, n = append([]string{"tasks", "--abs-time"}, append(plain, id)...), 4
	}
	r := run(Cmd{Name: "snap", Args: args})
	if strings.Contains(r.Stdout+r.Stderr, "no changes found") {
		return map[string]any{"changes": []Change{}}, nil
	}
	if r.Status != 0 || r.Error != "" {
		return nil, failure(Cmd{Name: "snap", Args: args}, r)
	}
	out := []Change{}
	for _, f := range columns(r.Stdout, n) {
		c := Change{}
		if n == 5 {
			c.ID, f = f[0], f[1:]
		}
		c.Status, c.Spawn, c.Ready, c.Summary = f[0], f[1], f[2], f[3]
		if c.Ready == "-" {
			c.Ready = ""
		}
		out = append(out, c)
	}
	if id != "" {
		return map[string]any{"id": id, "tasks": out}, nil
	}
	return map[string]any{"changes": out}, nil
}

// ActAnswer is what an act answers: the change snapd carries it out in.
type ActAnswer struct {
	Act    string `json:"act"`
	Snap   string `json:"snap,omitempty"`
	Change string `json:"change,omitempty"`
	Said   string `json:"said,omitempty"`
	Follow string `json:"follow,omitempty"`
}

// act runs one snap act with --no-wait, which answers snapd's change id at once.
func act(verb, name string, extra ...string) (ActAnswer, error) {
	args := append([]string{verb, "--no-wait"}, extra...)
	if name != "" {
		args = append(args, name)
	}
	r, err := call(Cmd{Name: "snap", Args: args, Root: true})
	if err != nil {
		return ActAnswer{}, err
	}
	out := ActAnswer{Act: verb, Snap: name}
	if id := strings.TrimSpace(r.Stdout); changeID.MatchString(id) {
		out.Change, out.Follow = id, "snapd_changes with id "+id
	} else {
		out.Said = strings.TrimSpace(r.Stdout + "\n" + r.Stderr)
	}
	return out, nil
}

var channelName = regexp.MustCompile(`^[a-z0-9][a-z0-9./_-]*$`)

// Install installs a snap.
func Install(name, channel string, classic bool) (ActAnswer, error) {
	extra := []string{}
	if channel != "" {
		if !channelName.MatchString(channel) {
			return ActAnswer{}, fmt.Errorf("%q is not a channel", channel)
		}
		extra = append(extra, "--channel="+channel)
	}
	if classic {
		extra = append(extra, "--classic")
	}
	return act("install", name, extra...)
}

// Remove removes a snap.
func Remove(name string, purge bool) (ActAnswer, error) {
	if purge {
		return act("remove", name, "--purge")
	}
	return act("remove", name)
}

// Refresh refreshes one snap, or all.
func Refresh(name string) (ActAnswer, error) {
	return act("refresh", name)
}