// Package agent implements the LLM tool-calling loop.
package agent

import (
	"context"
	"fmt"
	"sync"

	"github.com/crypto/byteclaw/internal/llm"
	"github.com/crypto/byteclaw/internal/store"
	"github.com/crypto/byteclaw/internal/tools"
)

// Event is streamed to clients while the agent works.
type Event struct {
	Type    string `json:"type"` // delta | tool_call | tool_result | done | error
	Text    string `json:"text,omitempty"`
	Tool    string `json:"tool,omitempty"`
	Args    string `json:"args,omitempty"`
	CallID  string `json:"call_id,omitempty"`
	Session string `json:"session_id,omitempty"`
}

type Agent struct {
	LLM           llm.Provider
	Tools         *tools.Registry
	Store         *store.Store
	SystemPrompt  string
	MaxIterations int
	HistoryLimit  int

	mu    sync.Mutex
	locks map[string]*sync.Mutex
}

func (a *Agent) sessionLock(id string) *sync.Mutex {
	a.mu.Lock()
	defer a.mu.Unlock()
	if a.locks == nil {
		a.locks = map[string]*sync.Mutex{}
	}
	l, ok := a.locks[id]
	if !ok {
		l = &sync.Mutex{}
		a.locks[id] = l
	}
	return l
}

// Run appends the user message to the session and loops until the model
// answers without tool calls. Returns the final assistant text.
func (a *Agent) Run(ctx context.Context, sessionID, userText string, emit func(Event)) (string, error) {
	if emit == nil {
		emit = func(Event) {}
	}
	l := a.sessionLock(sessionID)
	l.Lock()
	defer l.Unlock()

	if _, err := a.Store.AppendMessage(ctx, sessionID, llm.Message{Role: llm.RoleUser, Content: userText}); err != nil {
		return "", err
	}
	var specs []llm.ToolSpec
	if a.Tools != nil {
		specs = a.Tools.Specs()
	}
	for i := 0; i < a.MaxIterations; i++ {
		msgs, err := a.history(ctx, sessionID)
		if err != nil {
			return "", err
		}
		resp, err := a.LLM.Chat(ctx, llm.Request{Messages: msgs, Tools: specs}, func(t string) {
			emit(Event{Type: "delta", Text: t})
		})
		if err != nil {
			return "", err
		}
		asst := llm.Message{Role: llm.RoleAssistant, Content: resp.Content, ToolCalls: resp.ToolCalls}
		if _, err := a.Store.AppendMessage(ctx, sessionID, asst); err != nil {
			return "", err
		}
		if len(resp.ToolCalls) == 0 {
			return resp.Content, nil
		}
		for _, tc := range resp.ToolCalls {
			emit(Event{Type: "tool_call", Tool: tc.Name, Args: tc.Arguments, CallID: tc.ID})
			out := a.Tools.Call(ctx, tc.Name, tc.Arguments)
			emit(Event{Type: "tool_result", Tool: tc.Name, Text: clip(out, 2000), CallID: tc.ID})
			if _, err := a.Store.AppendMessage(ctx, sessionID, llm.Message{Role: llm.RoleTool, ToolCallID: tc.ID, Name: tc.Name, Content: out}); err != nil {
				return "", err
			}
		}
	}
	return "", fmt.Errorf("stopped after %d iterations without a final answer", a.MaxIterations)
}

func (a *Agent) history(ctx context.Context, sessionID string) ([]llm.Message, error) {
	stored, err := a.Store.Messages(ctx, sessionID, a.HistoryLimit)
	if err != nil {
		return nil, err
	}
	// A truncated window must not start with orphaned tool results.
	for len(stored) > 0 && stored[0].Role == llm.RoleTool {
		stored = stored[1:]
	}
	msgs := make([]llm.Message, 0, len(stored)+1)
	if a.SystemPrompt != "" {
		msgs = append(msgs, llm.Message{Role: llm.RoleSystem, Content: a.SystemPrompt})
	}
	for _, m := range stored {
		msgs = append(msgs, m.Message)
	}
	return msgs, nil
}

func clip(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "…"
}