package wing

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
)

// CompleteFunc is a chat completion, matching llm.Client.Complete. It must be
// backed by a local model: note content is sent to it.
type CompleteFunc func(ctx context.Context, system, user string) (string, error)

// maxNoteChars bounds what is sent for classification; the start of a note
// carries its subject, and a whole transcript is cost with no benefit.
const maxNoteChars = 6000

const classifySystem = `You file notes into a personal knowledge base. Pick the single wing the note belongs to, from the candidate list only.
Answer with ONLY a JSON object: {"wing": "<name or none>", "confidence": <0..1>, "reason": "<one short sentence>"}.
Use "none" when no candidate clearly fits. Do not invent wings.`

// LLMClassifier asks a model to pick a wing.
type LLMClassifier struct{ complete CompleteFunc }

// NewLLMClassifier wraps a chat completion as a Classifier.
func NewLLMClassifier(c CompleteFunc) *LLMClassifier { return &LLMClassifier{complete: c} }

type classifyReply struct {
	Wing       string  `json:"wing"`
	Confidence float64 `json:"confidence"`
	Reason     string  `json:"reason"`
}

// Classify implements Classifier. Any reply it cannot parse is an error, which
// the Resolver turns into an Exception: an unreadable answer is not a guess.
func (l *LLMClassifier) Classify(ctx context.Context, text string, candidates []string) (Suggestion, error) {
	if len(text) > maxNoteChars {
		text = text[:maxNoteChars]
	}
	user := "Candidate wings: " + strings.Join(candidates, ", ") + "\n\nNote:\n" + text
	out, err := l.complete(ctx, classifySystem, user)
	if err != nil {
		return Suggestion{}, fmt.Errorf("classify: %w", err)
	}
	raw, ok := firstJSONObject(out)
	if !ok {
		return Suggestion{}, errors.New("classify: reply contains no JSON object")
	}
	var r classifyReply
	if err := json.Unmarshal([]byte(raw), &r); err != nil {
		return Suggestion{}, fmt.Errorf("classify: parse reply: %w", err)
	}
	w := strings.TrimSpace(r.Wing)
	if strings.EqualFold(w, "none") {
		w = ""
	}
	return Suggestion{Wing: w, Confidence: clamp01(r.Confidence), Reason: r.Reason}, nil
}

// firstJSONObject returns the first balanced {...} in s, ignoring braces
// inside JSON strings, so prose or code fences around the reply are tolerated.
func firstJSONObject(s string) (string, bool) {
	start := strings.IndexByte(s, '{')
	if start < 0 {
		return "", false
	}
	depth, inStr, esc := 0, false, false
	for i := start; i < len(s); i++ {
		c := s[i]
		switch {
		case esc:
			esc = false
		case inStr && c == '\\':
			esc = true
		case c == '"':
			inStr = !inStr
		case inStr:
		case c == '{':
			depth++
		case c == '}':
			depth--
			if depth == 0 {
				return s[start : i+1], true
			}
		}
	}
	return "", false
}

func clamp01(f float64) float64 {
	switch {
	case f < 0:
		return 0
	case f > 1:
		return 1
	}
	return f
}