package inventory

import (
	"fmt"
	"os"
	"strings"

	"gopkg.in/yaml.v3"
)

// Inventory is the root document defining node classes and per-node specs.
type Inventory struct {
	Classes []NodeClass `yaml:"classes" json:"classes"`
	Nodes   []NodeSpec  `yaml:"nodes" json:"nodes"`
}

// NodeClass defines a reusable image/config for a group of nodes.
type NodeClass struct {
	Name           string            `yaml:"name" json:"name"`
	Arch           string            `yaml:"arch" json:"arch"`
	OS             string            `yaml:"os" json:"os"`
	Image          string            `yaml:"image" json:"image"`
	Checksum       string            `yaml:"checksum,omitempty" json:"checksum,omitempty"`
	K3sVersion     string            `yaml:"k3s_version,omitempty" json:"k3s_version,omitempty"`
	BootloaderNote string            `yaml:"bootloader_note,omitempty" json:"bootloader_note,omitempty"`
	DefaultLabels  map[string]string `yaml:"default_labels,omitempty" json:"default_labels,omitempty"`
	DefaultTaints  []string          `yaml:"default_taints,omitempty" json:"default_taints,omitempty"`
	CloudInit      string            `yaml:"cloud_init,omitempty" json:"cloud_init,omitempty"`
	BootOverlay    string            `yaml:"boot_overlay,omitempty" json:"boot_overlay,omitempty"` // path to overlay files for boot partition
	RootOverlay    string            `yaml:"root_overlay,omitempty" json:"root_overlay,omitempty"` // path to overlay files for rootfs
}

// NodeSpec captures per-node overrides and identity.
type NodeSpec struct {
	Name          string            `yaml:"name" json:"name"`
	Class         string            `yaml:"class" json:"class"`
	Hostname      string            `yaml:"hostname" json:"hostname"`
	IP            string            `yaml:"ip" json:"ip"`
	MAC           string            `yaml:"mac,omitempty" json:"mac,omitempty"`
	K3sRole       string            `yaml:"k3s_role" json:"k3s_role"`
	K3sVersion    string            `yaml:"k3s_version,omitempty" json:"k3s_version,omitempty"`
	K3sToken      string            `yaml:"k3s_token,omitempty" json:"k3s_token,omitempty"`
	K3sURL        string            `yaml:"k3s_url,omitempty" json:"k3s_url,omitempty"`
	Labels        map[string]string `yaml:"labels,omitempty" json:"labels,omitempty"`
	Taints        []string          `yaml:"taints,omitempty" json:"taints,omitempty"`
	LonghornDisks []LonghornDisk    `yaml:"longhorn_disks,omitempty" json:"longhorn_disks,omitempty"`
	USBScratch    *USBScratchDisk   `yaml:"usb_scratch,omitempty" json:"usb_scratch,omitempty"`
	SSHUser       string            `yaml:"ssh_user,omitempty" json:"ssh_user,omitempty"`
	SSHAuthorized []string          `yaml:"ssh_authorized_keys,omitempty" json:"ssh_authorized_keys,omitempty"`
	Notes         string            `yaml:"notes,omitempty" json:"notes,omitempty"`
}

// LonghornDisk describes an attached disk to mount for Longhorn.
type LonghornDisk struct {
	Mountpoint string `yaml:"mountpoint" json:"mountpoint"`
	UUID       string `yaml:"uuid" json:"uuid"`
	FS         string `yaml:"fs,omitempty" json:"fs,omitempty"`
}

// USBScratchDisk describes the recovery USB disk and its bind targets.
type USBScratchDisk struct {
	Mountpoint  string   `yaml:"mountpoint" json:"mountpoint"`
	UUID        string   `yaml:"uuid,omitempty" json:"uuid,omitempty"`
	Label       string   `yaml:"label,omitempty" json:"label,omitempty"`
	FS          string   `yaml:"fs,omitempty" json:"fs,omitempty"`
	BindTargets []string `yaml:"bind_targets,omitempty" json:"bind_targets,omitempty"`
}

// Load reads and parses an inventory file.
func Load(path string) (*Inventory, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("read inventory: %w", err)
	}
	var inv Inventory
	if err := yaml.Unmarshal(data, &inv); err != nil {
		return nil, fmt.Errorf("parse inventory: %w", err)
	}
	expandInventory(&inv)
	return &inv, nil
}

func expandInventory(inv *Inventory) {
	for idx := range inv.Classes {
		inv.Classes[idx].Name = os.ExpandEnv(inv.Classes[idx].Name)
		inv.Classes[idx].Arch = os.ExpandEnv(inv.Classes[idx].Arch)
		inv.Classes[idx].OS = os.ExpandEnv(inv.Classes[idx].OS)
		inv.Classes[idx].Image = os.ExpandEnv(inv.Classes[idx].Image)
		inv.Classes[idx].Checksum = os.ExpandEnv(inv.Classes[idx].Checksum)
		inv.Classes[idx].K3sVersion = os.ExpandEnv(inv.Classes[idx].K3sVersion)
		inv.Classes[idx].BootloaderNote = os.ExpandEnv(inv.Classes[idx].BootloaderNote)
		inv.Classes[idx].CloudInit = os.ExpandEnv(inv.Classes[idx].CloudInit)
		inv.Classes[idx].BootOverlay = os.ExpandEnv(inv.Classes[idx].BootOverlay)
		inv.Classes[idx].RootOverlay = os.ExpandEnv(inv.Classes[idx].RootOverlay)
		for key, value := range inv.Classes[idx].DefaultLabels {
			inv.Classes[idx].DefaultLabels[key] = os.ExpandEnv(value)
		}
		for taintIdx, value := range inv.Classes[idx].DefaultTaints {
			inv.Classes[idx].DefaultTaints[taintIdx] = os.ExpandEnv(value)
		}
	}
	for idx := range inv.Nodes {
		inv.Nodes[idx].Name = os.ExpandEnv(inv.Nodes[idx].Name)
		inv.Nodes[idx].Class = os.ExpandEnv(inv.Nodes[idx].Class)
		inv.Nodes[idx].Hostname = os.ExpandEnv(inv.Nodes[idx].Hostname)
		inv.Nodes[idx].IP = os.ExpandEnv(inv.Nodes[idx].IP)
		inv.Nodes[idx].MAC = os.ExpandEnv(inv.Nodes[idx].MAC)
		inv.Nodes[idx].K3sRole = os.ExpandEnv(inv.Nodes[idx].K3sRole)
		inv.Nodes[idx].K3sVersion = os.ExpandEnv(inv.Nodes[idx].K3sVersion)
		inv.Nodes[idx].K3sToken = os.ExpandEnv(inv.Nodes[idx].K3sToken)
		inv.Nodes[idx].K3sURL = os.ExpandEnv(inv.Nodes[idx].K3sURL)
		inv.Nodes[idx].SSHUser = os.ExpandEnv(inv.Nodes[idx].SSHUser)
		inv.Nodes[idx].Notes = os.ExpandEnv(inv.Nodes[idx].Notes)
		for key, value := range inv.Nodes[idx].Labels {
			inv.Nodes[idx].Labels[key] = os.ExpandEnv(value)
		}
		for taintIdx, value := range inv.Nodes[idx].Taints {
			inv.Nodes[idx].Taints[taintIdx] = os.ExpandEnv(value)
		}
		expandedKeys := make([]string, 0, len(inv.Nodes[idx].SSHAuthorized))
		for _, value := range inv.Nodes[idx].SSHAuthorized {
			expanded := strings.TrimSpace(os.ExpandEnv(value))
			if expanded == "" {
				continue
			}
			expandedKeys = append(expandedKeys, expanded)
		}
		inv.Nodes[idx].SSHAuthorized = expandedKeys
		for diskIdx := range inv.Nodes[idx].LonghornDisks {
			inv.Nodes[idx].LonghornDisks[diskIdx].Mountpoint = os.ExpandEnv(inv.Nodes[idx].LonghornDisks[diskIdx].Mountpoint)
			inv.Nodes[idx].LonghornDisks[diskIdx].UUID = os.ExpandEnv(inv.Nodes[idx].LonghornDisks[diskIdx].UUID)
			inv.Nodes[idx].LonghornDisks[diskIdx].FS = os.ExpandEnv(inv.Nodes[idx].LonghornDisks[diskIdx].FS)
		}
		if inv.Nodes[idx].USBScratch != nil {
			inv.Nodes[idx].USBScratch.Mountpoint = os.ExpandEnv(inv.Nodes[idx].USBScratch.Mountpoint)
			inv.Nodes[idx].USBScratch.UUID = os.ExpandEnv(inv.Nodes[idx].USBScratch.UUID)
			inv.Nodes[idx].USBScratch.Label = os.ExpandEnv(inv.Nodes[idx].USBScratch.Label)
			inv.Nodes[idx].USBScratch.FS = os.ExpandEnv(inv.Nodes[idx].USBScratch.FS)
			bindTargets := make([]string, 0, len(inv.Nodes[idx].USBScratch.BindTargets))
			for _, value := range inv.Nodes[idx].USBScratch.BindTargets {
				expanded := strings.TrimSpace(os.ExpandEnv(value))
				if expanded == "" {
					continue
				}
				bindTargets = append(bindTargets, expanded)
			}
			inv.Nodes[idx].USBScratch.BindTargets = bindTargets
		}
	}
}

// FindNode returns the node spec and class.
func (i *Inventory) FindNode(name string) (*NodeSpec, *NodeClass, error) {
	var node *NodeSpec
	for idx := range i.Nodes {
		if i.Nodes[idx].Name == name {
			node = &i.Nodes[idx]
			break
		}
	}
	if node == nil {
		return nil, nil, fmt.Errorf("node %s not found", name)
	}
	for idx := range i.Classes {
		if i.Classes[idx].Name == node.Class {
			return node, &i.Classes[idx], nil
		}
	}
	return node, nil, fmt.Errorf("class %s not found for node %s", node.Class, node.Name)
}
