package prompt

import (
    "errors"
    "io/fs"
    "os"
    "path/filepath"
    "strings"

    "gopkg.in/yaml.v3"
)

var (
    loadedRoles   map[string]RoleDefinition
    defaultInstr  string
)

// Initialize loads roles and default instructions from the configured directory.
// dir: base directory (e.g., /etc/chorus/prompts)
// defaultPath: optional explicit path to defaults file; if empty, will look for defaults.md or defaults.txt in dir.
func Initialize(dir string, defaultPath string) error {
    loadedRoles = make(map[string]RoleDefinition)

    // Load roles from all YAML files under dir
    if dir != "" {
        _ = filepath.WalkDir(dir, func(path string, d fs.DirEntry, err error) error {
            if err != nil || d == nil || d.IsDir() {
                return nil
            }
            name := d.Name()
            if strings.HasSuffix(name, ".yaml") || strings.HasSuffix(name, ".yml") {
                _ = loadRolesFile(path)
            }
            return nil
        })
    }

    // Load default instructions
    if defaultPath == "" && dir != "" {
        // Try defaults.md then defaults.txt in the directory
        tryPaths := []string{
            filepath.Join(dir, "defaults.md"),
            filepath.Join(dir, "defaults.txt"),
        }
        for _, p := range tryPaths {
            if b, err := os.ReadFile(p); err == nil {
                defaultInstr = string(b)
                break
            }
        }
    } else if defaultPath != "" {
        if b, err := os.ReadFile(defaultPath); err == nil {
            defaultInstr = string(b)
        }
    }
    return nil
}

func loadRolesFile(path string) error {
    data, err := os.ReadFile(path)
    if err != nil {
        return err
    }
    var rf RolesFile
    if err := yaml.Unmarshal(data, &rf); err != nil {
        return err
    }
    for id, def := range rf.Roles {
        def.ID = id
        loadedRoles[id] = def
    }
    return nil
}

// GetRole returns the role definition by ID if loaded.
func GetRole(id string) (RoleDefinition, bool) {
    r, ok := loadedRoles[id]
    return r, ok
}

// ListRoles returns IDs of loaded roles.
func ListRoles() []string {
    ids := make([]string, 0, len(loadedRoles))
    for id := range loadedRoles {
        ids = append(ids, id)
    }
    return ids
}

// GetDefaultInstructions returns the loaded default instructions (may be empty if not present on disk).
func GetDefaultInstructions() string {
    return defaultInstr
}

// ComposeSystemPrompt concatenates the role system prompt (S) with default instructions (D).
func ComposeSystemPrompt(roleID string) (string, error) {
    r, ok := GetRole(roleID)
    if !ok {
        return "", errors.New("role not found: " + roleID)
    }
    s := strings.TrimSpace(r.SystemPrompt)
    d := strings.TrimSpace(defaultInstr)
    switch {
    case s != "" && d != "":
        return s + "\n\n" + d, nil
    case s != "":
        return s, nil
    case d != "":
        return d, nil
    default:
        return "", nil
    }
}