// internal/skills/trainer/handlers.go
package trainer

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	iexec "github.com/mathiasbq/supervisor/internal/exec"
	"github.com/mathiasbq/supervisor/internal/session"
)

type trainArgs struct {
	SessionID string `json:"session_id"`
	Model     string `json:"model"`
}

// Handle dispatches the MCP tool call to the trainer handler.
func (s *Skill) Handle(ctx context.Context, tool string, args json.RawMessage) (json.RawMessage, error) {
	if tool != "trainer" {
		return nil, fmt.Errorf("unknown tool: %s", tool)
	}
	var a trainArgs
	if err := json.Unmarshal(args, &a); err != nil {
		return nil, fmt.Errorf("parse args: %w", err)
	}
	if a.SessionID == "" {
		return nil, fmt.Errorf("session_id is required")
	}
	if s.cfg.ExecutorFn == nil {
		return nil, fmt.Errorf("no executor configured")
	}

	model := a.Model
	if model == "" {
		model = s.cfg.DefaultModel
	}

	entries, err := session.Read(s.cfg.SessionsDir, a.SessionID)
	if err != nil {
		return nil, fmt.Errorf("read session log: %w", err)
	}

	// ── Step 1: Reader agent ─────────────────────────────────────────────────
	history := session.FormatHistory(entries, "")
	readerTask := fmt.Sprintf(
		"role: reader\nsession_id: %s\nbrain_dir: %s\n\n%s",
		a.SessionID, s.cfg.BrainDir, history,
	)
	readerResult, err := s.cfg.ExecutorFn(ctx, iexec.Request{
		SkillPrompt: s.cfg.ReaderPrompt,
		TaskPrompt:  readerTask,
		Model:       model,
		Tools:       "Read",
	})
	if err != nil {
		return nil, fmt.Errorf("reader agent: %w", err)
	}

	// ── Step 2: Writer agent (receives reader candidates) ────────────────────
	t0 := time.Now()
	writerTask := fmt.Sprintf(
		"role: writer\nsession_id: %s\nbrain_dir: %s\n\nreader_summary: %s\nreader_candidates:\n%s",
		a.SessionID, s.cfg.BrainDir, readerResult.Message, readerResult.RunnerOutput,
	)
	writerResult, err := s.cfg.ExecutorFn(ctx, iexec.Request{
		SkillPrompt: s.cfg.WriterPrompt,
		TaskPrompt:  writerTask,
		Model:       model,
		Tools:       "Read,Write",
	})
	if err != nil {
		return nil, fmt.Errorf("writer agent: %w", err)
	}

	_ = session.Append(s.cfg.SessionsDir, a.SessionID, session.Entry{
		SessionID:   a.SessionID,
		Timestamp:   time.Now(),
		Skill:       "trainer",
		Phase:       "trainer",
		Attempts:    session.AttemptsFrom(writerResult.Attempts),
		FinalStatus: writerResult.Status,
		ModelUsed:   writerResult.ModelUsed,
		DurationMs:  time.Since(t0).Milliseconds(),
		Message:     writerResult.Message,
	})

	b, err := json.Marshal(writerResult)
	if err != nil {
		return nil, fmt.Errorf("marshal result: %w", err)
	}
	return b, nil
}