vibetree/internal/assets/csv.go

package assets

import (
	"context"
	"encoding/csv"
	"fmt"
	"io"
	"sort"
	"strings"

	"inventree/internal/models"
)

const (
	colAssetID      = "asset id"
	colName         = "name"
	colItemType     = "item type"
	colFolder       = "folder"
	colStatus       = "status"
	colAssignedUser = "assigned user"
	colNotes        = "notes"
	cfColPrefix     = "cf_"
)

// ImportOptions controls CSV import behavior.
type ImportOptions struct {
	CreateMissingFolders bool
}

type ImportResult struct {
	Created int
	Updated int
	Errors  []string
}

// lookups bundles reference data needed to resolve CSV rows to foreign keys.
type lookups struct {
	itemTypesByName map[string]*models.ItemType
	statusesByName  map[string]int64
	usersByName     map[string]int64
	folderPathToID  map[string]int64
}

func normalize(s string) string { return strings.ToLower(strings.TrimSpace(s)) }

// Import reads assets from CSV data, creating or updating rows. Matching is
// done by Asset ID: if a row's Asset ID matches an existing asset, it is
// updated in place; otherwise a new asset is created (auto-generating an ID
// if the column is blank).
func (s *Store) Import(ctx context.Context, r io.Reader, opts ImportOptions,
	itemTypesList []models.ItemType, statusesList []models.Status, usersList []models.User) (*ImportResult, error) {

	lk := &lookups{
		itemTypesByName: map[string]*models.ItemType{},
		statusesByName:  map[string]int64{},
		usersByName:     map[string]int64{},
		folderPathToID:  map[string]int64{},
	}
	for i := range itemTypesList {
		lk.itemTypesByName[normalize(itemTypesList[i].Name)] = &itemTypesList[i]
	}
	for _, st := range statusesList {
		lk.statusesByName[normalize(st.Name)] = st.ID
	}
	for _, u := range usersList {
		lk.usersByName[normalize(u.Username)] = u.ID
	}
	allFolders, err := s.folders.All(ctx)
	if err != nil {
		return nil, err
	}
	byID := map[int64]models.Folder{}
	for _, f := range allFolders {
		byID[f.ID] = f
	}
	var pathOf func(id int64) string
	pathOf = func(id int64) string {
		f, ok := byID[id]
		if !ok {
			return ""
		}
		if f.ParentID == nil {
			return f.Name
		}
		return pathOf(*f.ParentID) + "/" + f.Name
	}
	for _, f := range allFolders {
		lk.folderPathToID[normalize(pathOf(f.ID))] = f.ID
	}

	cr := csv.NewReader(r)
	cr.FieldsPerRecord = -1
	header, err := cr.Read()
	if err != nil {
		if err == io.EOF {
			return &ImportResult{}, nil
		}
		return nil, fmt.Errorf("read CSV header: %w", err)
	}

	colIdx := map[string]int{}
	var cfCols []struct {
		key string
		idx int
	}
	for i, h := range header {
		n := normalize(h)
		if strings.HasPrefix(n, cfColPrefix) {
			cfCols = append(cfCols, struct {
				key string
				idx int
			}{key: strings.TrimPrefix(n, cfColPrefix), idx: i})
			continue
		}
		colIdx[n] = i
	}

	get := func(row []string, key string) string {
		idx, ok := colIdx[key]
		if !ok || idx >= len(row) {
			return ""
		}
		return strings.TrimSpace(row[idx])
	}

	result := &ImportResult{}
	rowNum := 1
	for {
		row, err := cr.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, fmt.Errorf("read CSV row %d: %w", rowNum, err)
		}
		rowNum++

		name := get(row, colName)
		if name == "" {
			result.Errors = append(result.Errors, fmt.Sprintf("row %d: name is required, skipped", rowNum))
			continue
		}
		itemTypeName := get(row, colItemType)
		it, ok := lk.itemTypesByName[normalize(itemTypeName)]
		if !ok {
			result.Errors = append(result.Errors, fmt.Sprintf("row %d: unknown item type %q, skipped", rowNum, itemTypeName))
			continue
		}
		statusName := get(row, colStatus)
		statusID, ok := lk.statusesByName[normalize(statusName)]
		if !ok {
			if len(statusesList) == 0 {
				result.Errors = append(result.Errors, fmt.Sprintf("row %d: no statuses configured, skipped", rowNum))
				continue
			}
			statusID = statusesList[0].ID
		}

		var folderID *int64
		folderPath := get(row, colFolder)
		if folderPath != "" {
			id, err := s.resolveFolderPath(ctx, lk, folderPath, opts.CreateMissingFolders)
			if err != nil {
				result.Errors = append(result.Errors, fmt.Sprintf("row %d: folder %q: %v", rowNum, folderPath, err))
				continue
			}
			folderID = &id
		}

		var assignedUserID *int64
		if uname := get(row, colAssignedUser); uname != "" {
			if id, ok := lk.usersByName[normalize(uname)]; ok {
				assignedUserID = &id
			} else {
				result.Errors = append(result.Errors, fmt.Sprintf("row %d: unknown assigned user %q, leaving unassigned", rowNum, uname))
			}
		}

		fieldValues := map[int64]string{}
		for _, cf := range cfCols {
			if cf.idx >= len(row) {
				continue
			}
			val := strings.TrimSpace(row[cf.idx])
			if val == "" {
				continue
			}
			for _, def := range it.Fields {
				if normalize(def.FieldKey) == cf.key {
					fieldValues[def.ID] = val
					break
				}
			}
		}

		asset := &models.Asset{
			AssetTag:       get(row, colAssetID),
			Name:           name,
			ItemTypeID:     it.ID,
			FolderID:       folderID,
			StatusID:       statusID,
			AssignedUserID: assignedUserID,
			Notes:          get(row, colNotes),
		}

		if asset.AssetTag != "" {
			existing, err := s.GetByTag(ctx, asset.AssetTag)
			if err == nil {
				asset.ID = existing.ID
				if err := s.Update(ctx, asset, fieldValues); err != nil {
					result.Errors = append(result.Errors, fmt.Sprintf("row %d: update failed: %v", rowNum, err))
					continue
				}
				result.Updated++
				continue
			}
		}

		if err := s.Create(ctx, asset, fieldValues); err != nil {
			result.Errors = append(result.Errors, fmt.Sprintf("row %d: create failed: %v", rowNum, err))
			continue
		}
		result.Created++
	}

	return result, nil
}

func (s *Store) resolveFolderPath(ctx context.Context, lk *lookups, path string, createMissing bool) (int64, error) {
	norm := normalize(path)
	if id, ok := lk.folderPathToID[norm]; ok {
		return id, nil
	}
	if !createMissing {
		return 0, fmt.Errorf("folder not found")
	}

	segments := strings.Split(path, "/")
	var parentID *int64
	built := ""
	for _, seg := range segments {
		seg = strings.TrimSpace(seg)
		if seg == "" {
			continue
		}
		if built == "" {
			built = seg
		} else {
			built = built + "/" + seg
		}
		if id, ok := lk.folderPathToID[normalize(built)]; ok {
			parentID = &id
			continue
		}
		f := &models.Folder{Name: seg, ParentID: parentID}
		if err := s.folders.Create(ctx, f); err != nil {
			return 0, err
		}
		lk.folderPathToID[normalize(built)] = f.ID
		parentID = &f.ID
	}
	if parentID == nil {
		return 0, fmt.Errorf("empty folder path")
	}
	return *parentID, nil
}

// Export writes all assets (optionally filtered) to CSV, with one column per
// distinct custom field key used across all item types.
func (s *Store) Export(ctx context.Context, w io.Writer, f ListFilter) error {
	fieldDefs, err := s.allFieldKeys(ctx)
	if err != nil {
		return err
	}

	cw := csv.NewWriter(w)
	header := []string{"Asset ID", "Name", "Item Type", "Folder", "Status", "Assigned User", "Notes", "Created At", "Updated At"}
	for _, k := range fieldDefs {
		header = append(header, cfColPrefix+k)
	}
	if err := cw.Write(header); err != nil {
		return err
	}

	f.Page = 1
	f.PageSize = 1000
	for {
		batch, total, err := s.List(ctx, f)
		if err != nil {
			return err
		}
		for i := range batch {
			a, err := s.Get(ctx, batch[i].ID)
			if err != nil {
				return err
			}
			row := []string{
				a.AssetTag, a.Name, a.ItemTypeName, folderPathString(a.FolderPath), a.StatusName, a.AssignedUsername, a.Notes,
				a.CreatedAt.Format("2006-01-02 15:04:05"), a.UpdatedAt.Format("2006-01-02 15:04:05"),
			}
			values := map[string]string{}
			for _, fv := range a.CustomFields {
				values[fv.Def.FieldKey] = fv.Value
			}
			for _, k := range fieldDefs {
				row = append(row, values[k])
			}
			if err := cw.Write(row); err != nil {
				return err
			}
		}
		if f.Page*f.PageSize >= total || len(batch) == 0 {
			break
		}
		f.Page++
	}
	cw.Flush()
	return cw.Error()
}

func folderPathString(path []models.BreadcrumbItem) string {
	names := make([]string, len(path))
	for i, p := range path {
		names[i] = p.Name
	}
	return strings.Join(names, "/")
}

func (s *Store) allFieldKeys(ctx context.Context) ([]string, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT DISTINCT field_key FROM custom_field_defs`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	seen := map[string]bool{}
	var out []string
	for rows.Next() {
		var k string
		if err := rows.Scan(&k); err != nil {
			return nil, err
		}
		if !seen[k] {
			seen[k] = true
			out = append(out, k)
		}
	}
	sort.Strings(out)
	return out, rows.Err()
}