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

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

	"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.CompleteFunc == 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 ────────────────────────────────────────────────────────
	history := session.FormatHistory(entries, "")
	readerTask := fmt.Sprintf(
		"role: reader\nsession_id: %s\nbrain_dir: %s\n\n%s",
		a.SessionID, s.cfg.BrainDir, history,
	)
	readerText, _, err := s.cfg.CompleteFunc(ctx, model, s.cfg.ReaderPrompt, readerTask)
	if err != nil {
		return nil, fmt.Errorf("reader: %w", err)
	}

	// ── Step 2: Writer (receives reader output) ───────────────────────────────
	t0 := time.Now()
	writerTask := fmt.Sprintf(
		"role: writer\nsession_id: %s\nbrain_dir: %s\n\nreader_analysis:\n%s",
		a.SessionID, s.cfg.BrainDir, readerText,
	)
	writerText, dur, err := s.cfg.CompleteFunc(ctx, model, s.cfg.WriterPrompt, writerTask)
	if err != nil {
		return nil, fmt.Errorf("writer: %w", err)
	}

	msg := writerText
	if len(msg) > 200 {
		msg = msg[:200]
	}
	_ = session.Append(s.cfg.SessionsDir, a.SessionID, session.Entry{
		SessionID:   a.SessionID,
		Timestamp:   time.Now(),
		Skill:       "trainer",
		Phase:       "trainer",
		FinalStatus: "ok",
		ModelUsed:   model,
		DurationMs:  time.Since(t0).Milliseconds(),
		Message:     msg,
	})

	return json.Marshal(map[string]any{
		"reader_analysis": readerText,
		"writer_output":   writerText,
		"model":           model,
		"duration_ms":     dur,
	})
}