package agent

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/sdaduanbilei/agent-team/internal/llm"
	"gopkg.in/yaml.v3"
)

type Config struct {
	Name        string   `yaml:"name"`
	Description string   `yaml:"description"`
	Provider    string   `yaml:"provider"`
	Model       string   `yaml:"model"`
	BaseURL     string   `yaml:"base_url"`
	APIKeyEnv   string   `yaml:"api_key_env"`
	Skills      []string `yaml:"skills"`
}

type Agent struct {
	Config Config
	Soul   string // system prompt from SOUL.md
	Dir    string // agents/<name>/
	client *llm.Client
}

func Load(dir string) (*Agent, error) {
	agentMD, err := os.ReadFile(filepath.Join(dir, "AGENT.md"))
	if err != nil {
		return nil, err
	}
	cfg, err := parseFrontmatter(agentMD)
	if err != nil {
		return nil, fmt.Errorf("parse AGENT.md: %w", err)
	}
	soul, _ := os.ReadFile(filepath.Join(dir, "SOUL.md"))

	if cfg.Provider == "" {
		cfg.Provider = "deepseek"
	}
	if cfg.APIKeyEnv == "" {
		cfg.APIKeyEnv = "DEEPSEEK_API_KEY"
	}

	client, err := llm.New(cfg.Provider, cfg.Model, cfg.BaseURL, cfg.APIKeyEnv)
	if err != nil {
		return nil, err
	}
	return &Agent{Config: cfg, Soul: string(soul), Dir: dir, client: client}, nil
}

func parseFrontmatter(data []byte) (Config, error) {
	var cfg Config
	if !bytes.HasPrefix(data, []byte("---")) {
		return cfg, fmt.Errorf("missing frontmatter")
	}
	parts := bytes.SplitN(data, []byte("---"), 3)
	if len(parts) < 3 {
		return cfg, fmt.Errorf("invalid frontmatter")
	}
	return cfg, yaml.Unmarshal(parts[1], &cfg)
}

// Memory returns concatenated memory files content.
func (a *Agent) Memory() string {
	memDir := filepath.Join(a.Dir, "memory")
	entries, err := os.ReadDir(memDir)
	if err != nil {
		return ""
	}
	var sb strings.Builder
	for _, e := range entries {
		if strings.HasSuffix(e.Name(), ".md") {
			data, _ := os.ReadFile(filepath.Join(memDir, e.Name()))
			sb.Write(data)
			sb.WriteString("\n")
		}
	}
	return sb.String()
}

// SaveMemory appends/updates a memory file.
func (a *Agent) SaveMemory(filename, content string) error {
	memDir := filepath.Join(a.Dir, "memory")
	os.MkdirAll(memDir, 0755)
	return os.WriteFile(filepath.Join(memDir, filename), []byte(content), 0644)
}

// Chat sends messages and streams response tokens via onToken callback.
func (a *Agent) Chat(ctx context.Context, msgs []llm.Message, onToken func(string)) (string, error) {
	return a.client.Stream(ctx, msgs, onToken)
}

// BuildSystemPrompt constructs the full system prompt with soul + memory + injected context.
func (a *Agent) BuildSystemPrompt(extraContext string) string {
	var sb strings.Builder
	sb.WriteString(a.Soul)
	if mem := a.Memory(); mem != "" {
		sb.WriteString("\n\n<memory>\n")
		sb.WriteString(mem)
		sb.WriteString("</memory>")
	}
	if extraContext != "" {
		sb.WriteString("\n\n")
		sb.WriteString(extraContext)
	}
	return sb.String()
}