// Package persona manages AI character loading and system prompt injection.
package persona

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
)

// Persona defines an AI character's configuration.
type Persona struct {
	Name         string   `yaml:"name" json:"name"`
	Description  string   `yaml:"description" json:"description"`
	Greeting     string   `yaml:"greeting" json:"greeting"`
	Traits       []string `yaml:"traits" json:"traits"`
	Style        string   `yaml:"style" json:"style"`
	Constraints  []string `yaml:"constraints" json:"constraints"`

	// IoT-specific personality
	IoTResponseStyle string `yaml:"iot_response_style" json:"iot_response_style,omitempty"`
}

// Loader loads persona configurations from a directory.
type Loader struct {
	personas map[string]*Persona
}

// NewLoader loads all YAML persona files from the given directory.
func NewLoader(dir string) (*Loader, error) {
	l := &Loader{personas: make(map[string]*Persona)}

	entries, err := os.ReadDir(dir)
	if err != nil {
		// Return empty loader if directory doesn't exist
		if os.IsNotExist(err) {
			return l, nil
		}
		return nil, fmt.Errorf("read persona dir %s: %w", dir, err)
	}

	for _, entry := range entries {
		if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".yaml") && !strings.HasSuffix(entry.Name(), ".yml") {
			continue
		}
		path := filepath.Join(dir, entry.Name())
		p, err := loadFile(path)
		if err != nil {
			return nil, fmt.Errorf("load persona %s: %w", entry.Name(), err)
		}
		name := strings.TrimSuffix(entry.Name(), filepath.Ext(entry.Name()))
		p.Name = name
		l.personas[name] = p
	}

	return l, nil
}

// Get returns a persona by name, or nil if not found.
func (l *Loader) Get(name string) *Persona {
	return l.personas[name]
}

// List returns all loaded persona names.
func (l *Loader) List() []string {
	names := make([]string, 0, len(l.personas))
	for name := range l.personas {
		names = append(names, name)
	}
	return names
}

// Count returns the number of loaded personas.
func (l *Loader) Count() int {
	return len(l.personas)
}

// Default returns the first persona, or a built-in default.
func (l *Loader) Default() *Persona {
	for _, p := range l.personas {
		return p
	}
	return BuiltinDefault()
}

// loadFile reads a single persona YAML file.
func loadFile(path string) (*Persona, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	return parseYAML(data)
}

// parseYAML parses persona YAML content (simple key-value parser).
func parseYAML(data []byte) (*Persona, error) {
	p := &Persona{}
	lines := strings.Split(string(data), "\n")
	var currentList *[]string

	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}

		if strings.HasPrefix(line, "- ") && currentList != nil {
			*currentList = append(*currentList, strings.TrimPrefix(line, "- "))
			continue
		}
		currentList = nil

		parts := strings.SplitN(line, ":", 2)
		if len(parts) != 2 {
			continue
		}
		key := strings.TrimSpace(parts[0])
		val := strings.TrimSpace(parts[1])

		switch key {
		case "name":
			p.Name = val
		case "description":
			p.Description = val
		case "greeting":
			p.Greeting = val
		case "style":
			p.Style = val
		case "iot_response_style":
			p.IoTResponseStyle = val
		case "traits":
			currentList = &p.Traits
		case "constraints":
			currentList = &p.Constraints
		}
	}
	return p, nil
}

// BuiltinDefault returns a hardcoded default persona.
func BuiltinDefault() *Persona {
	return &Persona{
		Name:        "default",
		Description: "一个友好的 AI 助手",
		Greeting:    "你好呀~ 有什么可以帮你的吗？",
		Traits:      []string{"友善", "细心", "幽默"},
		Style:       "用轻松自然的中文交流，偶尔加一点俏皮表情",
		Constraints: []string{"不编造事实", "不提供医疗法律建议", "保持礼貌"},
	}
}

// Injector builds the system prompt from a persona.
type Injector struct {
	persona *Persona
}

// NewInjector creates a system prompt injector for a persona.
func NewInjector(p *Persona) *Injector {
	if p == nil {
		p = BuiltinDefault()
	}
	return &Injector{persona: p}
}

// BuildSystemPrompt constructs the full system prompt string.
func (inj *Injector) BuildSystemPrompt() string {
	var sb strings.Builder

	sb.WriteString(fmt.Sprintf("你是%s，%s。\n", inj.persona.Name, inj.persona.Description))

	if len(inj.persona.Traits) > 0 {
		sb.WriteString(fmt.Sprintf("性格特点: %s\n", strings.Join(inj.persona.Traits, "、")))
	}
	if inj.persona.Style != "" {
		sb.WriteString(fmt.Sprintf("说话风格: %s\n", inj.persona.Style))
	}
	if len(inj.persona.Constraints) > 0 {
		sb.WriteString("遵守规则:\n")
		for _, c := range inj.persona.Constraints {
			sb.WriteString(fmt.Sprintf("- %s\n", c))
		}
	}

	return sb.String()
}