kb/client/cmd/add.go

package cmd

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/kb-search/kb/internal/api"
	"github.com/kb-search/kb/internal/output"
	"github.com/spf13/cobra"
)

type uploadResult struct {
	Raw       interface{}
	Duplicate bool
	DocID     float64
	JobID     float64
	Title     string
}

func (r *uploadResult) duplicateMsg() string {
	if r.DocID > 0 {
		return fmt.Sprintf("Already imported: %s (doc ID: %.0f)", r.Title, r.DocID)
	}
	return fmt.Sprintf("Already queued: %s (job ID: %.0f)", r.Title, r.JobID)
}

var supportedExts = map[string]bool{
	".pdf":  true,
	".docx": true,
	".html": true,
	".md":   true,
	".txt":  true,
	".py":   true,
	".sh":   true,
	".go":   true,
}

var addfileCmd = &cobra.Command{
	Use:   "addfile <path>",
	Short: "Upload a file or directory to the knowledge base",
	Args:  cobra.ExactArgs(1),
	RunE:  runAddfile,
}

func init() {
	addfileCmd.Flags().String("tags", "", "tags (comma-separated)")
	addfileCmd.Flags().BoolP("recursive", "r", false, "recursively add directory contents")
	rootCmd.AddCommand(addfileCmd)
}

func runAddfile(cmd *cobra.Command, args []string) error {
	tags, _ := cmd.Flags().GetString("tags")
	recursive, _ := cmd.Flags().GetBool("recursive")

	client := api.NewClient()

	path := args[0]
	info, err := os.Stat(path)
	if err != nil {
		return fmt.Errorf("cannot access %s: %w", path, err)
	}

	if !info.IsDir() {
		// Validate extension
		ext := strings.ToLower(filepath.Ext(path))
		if !supportedExts[ext] {
			supported := make([]string, 0, len(supportedExts))
			for e := range supportedExts {
				supported = append(supported, e)
			}
			sort.Strings(supported)
			return fmt.Errorf("unsupported file type %q — supported: %s", ext, strings.Join(supported, ", "))
		}

		// Single file upload
		result, err := uploadFile(client, path, tags)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}

		if output.IsJSON() {
			output.PrintJSON([]interface{}{result.Raw})
		} else if result.Duplicate {
			fmt.Println(result.duplicateMsg())
		} else {
			fmt.Printf("Queued: %s\n", filepath.Base(path))
		}
		return nil
	}

	// Directory upload
	var files []string
	walkFn := func(p string, d os.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if d.IsDir() {
			if !recursive && p != path {
				return filepath.SkipDir
			}
			return nil
		}
		ext := strings.ToLower(filepath.Ext(p))
		if supportedExts[ext] {
			files = append(files, p)
		}
		return nil
	}

	if err := filepath.WalkDir(path, walkFn); err != nil {
		return fmt.Errorf("failed to walk directory: %w", err)
	}

	var results []interface{}
	queued := 0
	duplicates := 0
	for _, f := range files {
		result, err := uploadFile(client, f, tags)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error uploading %s: %v\n", f, err)
			continue
		}
		results = append(results, result.Raw)
		if result.Duplicate {
			duplicates++
			if !output.IsJSON() {
				fmt.Println(result.duplicateMsg())
			}
		} else {
			queued++
			if !output.IsJSON() {
				fmt.Printf("Queued: %s\n", filepath.Base(f))
			}
		}
	}

	if output.IsJSON() {
		output.PrintJSON(results)
	} else if duplicates > 0 {
		fmt.Printf("Queued: %d files, %d duplicates skipped\n", queued, duplicates)
	} else {
		fmt.Printf("Queued: %d files\n", queued)
	}
	return nil
}

func uploadFile(client *api.Client, path, tags string) (*uploadResult, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, fmt.Errorf("cannot open %s: %w", path, err)
	}
	defer f.Close()

	fields := make(map[string]string)
	if tags != "" {
		fields["tags"] = tags
	}

	upload := &api.FileUpload{
		FieldName: "file",
		FileName:  filepath.Base(path),
		Reader:    f,
	}

	resp, err := client.PostMultipart("/api/v1/jobs", fields, upload)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode == http.StatusConflict {
		var raw json.RawMessage
		if err := api.DecodeJSON(resp, &raw); err != nil {
			return nil, fmt.Errorf("failed to decode response: %w", err)
		}
		var dupResp struct {
			DocumentID float64 `json:"document_id"`
			JobID      float64 `json:"job_id"`
			Title      string  `json:"title"`
		}
		json.Unmarshal(raw, &dupResp)
		var rawIface interface{}
		json.Unmarshal(raw, &rawIface)
		return &uploadResult{
			Raw:       rawIface,
			Duplicate: true,
			DocID:     dupResp.DocumentID,
			JobID:     dupResp.JobID,
			Title:     dupResp.Title,
		}, nil
	}

	if err := api.CheckError(resp); err != nil {
		return nil, err
	}

	var result interface{}
	if err := api.DecodeJSON(resp, &result); err != nil {
		return nil, fmt.Errorf("failed to decode response: %w", err)
	}
	return &uploadResult{Raw: result}, nil
}