package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"strconv"
	"time"
)

// Jev (TypeSafe) is a classifier, not a text model: it returns typed choices
// with calibrated confidence. The Sentry uses it for triage, and escalates to
// the LLM when Jev is unsure, when an item needs subtasks written, or when Jev
// fails. Docs: https://docs.typesafe.ai/introduction/quickstart

type JevClient struct {
	url, key, model string
	threshold       float64 // both answers must reach this confidence to stand alone
	client          *http.Client
}

// JevCall is one classification: a label and a confidence per question.
type JevCall struct {
	Mode, TaskType         string
	ModeConf, TaskTypeConf float64
	Model                  string
	Tokens                 int
}

// NewJev returns nil (no classifier) unless TYPESAFE_API_KEY is set.
//
//	TYPESAFE_API_KEY         the key from the TypeSafe console
//	CROWS_JEV_MODEL          default jev-latest
//	CROWS_JEV_THRESHOLD      default 0.7
func NewJev() *JevClient {
	key := os.Getenv("TYPESAFE_API_KEY")
	if key == "" {
		return nil
	}
	model := os.Getenv("CROWS_JEV_MODEL")
	if model == "" {
		model = "jev-latest"
	}
	threshold := 0.7
	if v, err := strconv.ParseFloat(os.Getenv("CROWS_JEV_THRESHOLD"), 64); err == nil && v > 0 && v <= 1 {
		threshold = v
	}
	return &JevClient{url: "https://api.typesafe.ai/v1/systemone", key: key, model: model, threshold: threshold, client: &http.Client{}}
}

func (j *JevClient) Name() string {
	return fmt.Sprintf("Jev (%s, escalates below %.0f%%)", j.model, j.threshold*100)
}

var jevQuestions = map[string]any{
	"mode": map[string]any{
		"type":         "choice",
		"instructions": "How should a team of AI agents handle this todo item?",
		"criteria": map[string]string{
			"solo":   "Simple or single-step: one agent can answer it directly.",
			"split":  "Large, with 2 to 4 distinct parts that can be done independently and in parallel.",
			"murder": "Ambiguous, high-stakes or a judgment call: several agents should answer independently, then reconcile.",
			"other":  "None of the above fits well.",
		},
	},
	"task_type": map[string]any{
		"type":         "choice",
		"instructions": "What kind of task is this todo item?",
		"criteria": map[string]string{
			"retrieval": "Looking up or recalling facts.",
			"writing":   "Drafting or editing text.",
			"analysis":  "Comparing, evaluating, deciding or reasoning about options.",
			"other":     "Anything else.",
		},
	},
}

// Answer is one Jev verdict, any question type. Pointed fields are present only
// for their type: choice sets Choice/Confidence, score sets Score/Confidence,
// noul sets Noul (noul has no confidence — see JEV_STRATEGIES.md).
type Answer struct {
	Type          string             `json:"type"`
	Choice        string             `json:"choice,omitempty"`
	Confidence    *float64           `json:"confidence,omitempty"`
	Score         *float64           `json:"score,omitempty"`
	Noul          *float64           `json:"noul,omitempty"`
	Probabilities map[string]float64 `json:"probabilities,omitempty"`
	Legend        map[string]any     `json:"legend,omitempty"`
}

// SystemOne POSTs one state and a set of questions, mixing types freely, and
// returns the answers keyed by question name. It is the single network path for
// every Jev call (Classify and the probes), so auth, status and usage handling
// stays in one place. Returns total tokens and the model that answered.
func (j *JevClient) SystemOne(ctx context.Context, state any, questions map[string]any) (map[string]Answer, int, string, error) {
	body, _ := json.Marshal(map[string]any{"model": j.model, "state": state, "questions": questions})
	req, err := http.NewRequestWithContext(ctx, "POST", j.url, bytes.NewReader(body))
	if err != nil {
		return nil, 0, "", err
	}
	req.Header.Set("Authorization", "Bearer "+j.key)
	req.Header.Set("Content-Type", "application/json")
	resp, err := j.client.Do(req)
	if err != nil {
		return nil, 0, "", err
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if resp.StatusCode != http.StatusOK {
		return nil, 0, "", fmt.Errorf("Jev returned %d: %s", resp.StatusCode, clip(raw, 300))
	}
	var out struct {
		Model   string            `json:"model"`
		Answers map[string]Answer `json:"answers"`
		Usage   struct {
			InputTokens  int `json:"input_tokens"`
			OutputTokens int `json:"output_tokens"`
		} `json:"usage"`
	}
	if err := json.Unmarshal(raw, &out); err != nil {
		return nil, 0, "", fmt.Errorf("Jev sent an unreadable reply: %w", err)
	}
	model := out.Model
	if model == "" {
		model = j.model
	}
	return out.Answers, out.Usage.InputTokens + out.Usage.OutputTokens, model, nil
}

func (j *JevClient) Classify(ctx context.Context, item string) (JevCall, error) {
	answers, tokens, model, err := j.SystemOne(ctx, item, jevQuestions)
	if err != nil {
		return JevCall{}, err
	}
	mode, task := answers["mode"], answers["task_type"]
	if mode.Choice == "" || task.Choice == "" {
		return JevCall{}, errors.New("Jev's reply was missing an answer")
	}
	return JevCall{
		Mode: mode.Choice, TaskType: task.Choice,
		ModeConf: deref(mode.Confidence), TaskTypeConf: deref(task.Confidence),
		Model: model, Tokens: tokens,
	}, nil
}

func deref(f *float64) float64 {
	if f == nil {
		return 0
	}
	return *f
}

// modeQuestions classifies how a request should be taken up (Epic C). The
// escape option "other" keeps Jev from forcing a wrong class; mode is discrete,
// so choice, not score (JEV_STRATEGIES.md strategy 3).
var modeQuestions = map[string]any{
	"mode": map[string]any{
		"type":         "choice",
		"instructions": "How should this chatroom request be taken up?",
		"criteria": map[string]string{
			"solo":       "Simple: one crow answers it directly.",
			"conspiracy": "A judgment call worth several distinct readings: several crows take the same request from their different strengths.",
			"split":      "Large with distinct parts: separate subtasks, one owner each.",
			"other":      "None of these fits well.",
		},
	},
}

// ClassifyMode returns the mode when Jev is set and its choice is at or above
// the band. "other" reads as unsure (not ok) so callers fall back to a header.
func (j *JevClient) ClassifyMode(ctx context.Context, text string) (string, map[string]float64, bool) {
	if j == nil {
		return "", nil, false
	}
	answers, _, _, err := j.SystemOne(ctx, text, modeQuestions)
	if err != nil {
		return "", nil, false
	}
	a := answers["mode"]
	if a.Choice == "" || a.Choice == "other" || a.Confidence == nil || *a.Confidence < j.threshold {
		return "", nil, false
	}
	return a.Choice, a.Probabilities, true
}

// ClassifyMode classifies a request under the same guardrails as any call:
// flock cap, timeout and flight log, with the probabilities recorded. Returns
// the mode and ok; below-band or absent Jev means not ok.
func (f *Flock) ClassifyMode(ctx context.Context, text string, budgets ...*Budget) (string, bool) {
	if f.jev == nil {
		return "", false
	}
	for _, b := range budgets {
		if err := b.Check(); err != nil {
			return "", false
		}
	}
	f.sem <- struct{}{}
	defer func() { <-f.sem }()
	cctx, cancel := context.WithTimeout(ctx, f.g.CrowTimeout)
	defer cancel()
	start := time.Now()
	mode, probs, ok := f.jev.ClassifyMode(cctx, text)
	e := Event{Kind: "mode", Item: text, Crow: "Sentry", Model: f.jev.model, Ms: time.Since(start).Milliseconds(), Prob: probs, Mode: mode}
	if f.log != nil {
		f.log.Write(e)
	}
	return mode, ok
}

func (c JevCall) confident(threshold float64) bool {
	return c.ModeConf >= threshold && c.TaskTypeConf >= threshold
}

// kindQuestions tags a nest message with its kind (A1: Jev first). The escape
// option "say" keeps Jev from forcing a wrong label: any message it isn't sure
// about falls back to the safest kind.
var kindQuestions = map[string]any{
	"kind": map[string]any{
		"type":         "choice",
		"instructions": "What kind of chatroom message is this?",
		"criteria": map[string]string{
			"request": "Asks a member or the room for work.",
			"claim":   "Takes ownership of a request.",
			"deliver": "Posts the result of a request.",
			"refer":   "Points to a file or artifact instead of pasting it.",
			"say":     "Ordinary talk, or unclear.",
		},
	},
}

// ClassifyKind returns the kind when Jev is set and its choice sits at or
// above the confidence band, with the choice probabilities for the log. Not ok
// means the caller falls back to the poster's header (EVENTS.md tagging rule).
func (j *JevClient) ClassifyKind(ctx context.Context, text string) (string, map[string]float64, bool) {
	if j == nil {
		return "", nil, false
	}
	answers, _, _, err := j.SystemOne(ctx, text, kindQuestions)
	if err != nil {
		return "", nil, false
	}
	a := answers["kind"]
	if a.Choice == "" || a.Confidence == nil || *a.Confidence < j.threshold {
		return "", nil, false
	}
	return a.Choice, a.Probabilities, true
}

// divergenceQuestions asks whether a message restates an earlier point. It is
// a noul: JEV_STRATEGIES strategy 3 (crisp yes/no) and noul has no confidence,
// so the probability itself is thresholded (band = threshold).
var divergenceQuestions = map[string]any{
	"restates_earlier": map[string]any{
		"type":         "noul",
		"instructions": "Does `this` make a point already made in `earlier`, adding no new angle?",
		"criteria": map[string]string{
			"true":  "It repeats an earlier point without a genuinely new angle.",
			"false": "It adds something different from what `earlier` already said.",
		},
	},
}

// ClassifyDivergence returns the probability that a message restates earlier
// ones, and ok when a value came back (noul always carries one; ok is false
// only on transport error or a missing answer).
func (j *JevClient) ClassifyDivergence(ctx context.Context, state any) (float64, bool) {
	if j == nil {
		return 0, false
	}
	answers, _, _, err := j.SystemOne(ctx, state, divergenceQuestions)
	if err != nil {
		return 0, false
	}
	a := answers["restates_earlier"]
	if a.Noul == nil {
		return 0, false
	}
	return *a.Noul, true
}

// ClassifyDivergence runs the restatement check under the flock guardrails
// (cap, timeout, flight log with the probability recorded).
func (f *Flock) ClassifyDivergence(ctx context.Context, state any, budgets ...*Budget) (float64, bool) {
	if f.jev == nil {
		return 0, false
	}
	for _, b := range budgets {
		if err := b.Check(); err != nil {
			return 0, false
		}
	}
	f.sem <- struct{}{}
	defer func() { <-f.sem }()
	cctx, cancel := context.WithTimeout(ctx, f.g.CrowTimeout)
	defer cancel()
	start := time.Now()
	p, ok := f.jev.ClassifyDivergence(cctx, state)
	if f.log != nil {
		f.log.Write(Event{Kind: "divergence", Crow: "Sentry", Model: f.jev.model, Ms: time.Since(start).Milliseconds(), Prob: map[string]float64{"restates_earlier": p}, Note: fmt.Sprint(ok && p >= f.jev.threshold)})
	}
	return p, ok
}

// ClassifyKind tags a nest message under the same guardrails as any call:
// flock cap, timeout and flight log, with the probabilities recorded.
func (f *Flock) ClassifyKind(ctx context.Context, text string, budgets ...*Budget) (string, map[string]float64, bool) {
	if f.jev == nil {
		return "", nil, false
	}
	for _, b := range budgets {
		if err := b.Check(); err != nil {
			return "", nil, false
		}
	}
	f.sem <- struct{}{}
	defer func() { <-f.sem }()
	cctx, cancel := context.WithTimeout(ctx, f.g.CrowTimeout)
	defer cancel()
	start := time.Now()
	kind, probs, ok := f.jev.ClassifyKind(cctx, text)
	e := Event{Kind: "tag", Item: text, Crow: "Sentry", Model: f.jev.model, Ms: time.Since(start).Milliseconds(), Prob: probs, Note: kind}
	if !ok {
		e.Note = "" // unsure: the header decides
	}
	if f.log != nil {
		f.log.Write(e)
	}
	return kind, probs, ok
}

func (c JevCall) summary() string {
	return fmt.Sprintf("Jev: %s %.0f%%, %s %.0f%%", c.Mode, c.ModeConf*100, c.TaskType, c.TaskTypeConf*100)
}

// classify runs Jev under the same guardrails as any crow: budget, flock cap,
// timeout, flight log, and flying/landed events for the Sentry.
func (f *Flock) classify(ctx context.Context, w Watcher, item string, budgets ...*Budget) (JevCall, error) {
	for _, b := range budgets {
		if err := b.Check(); err != nil {
			return JevCall{}, err
		}
	}
	f.sem <- struct{}{}
	defer func() { <-f.sem }()
	if w != nil {
		w.Flying("Sentry")
	}
	cctx, cancel := context.WithTimeout(ctx, f.g.CrowTimeout)
	defer cancel()
	start := time.Now()
	call, err := f.jev.Classify(cctx, item)
	for _, b := range budgets {
		b.Add(call.Tokens)
	}
	e := Event{Kind: "flight", Crow: "Sentry", Model: call.Model, Tokens: call.Tokens, Ms: time.Since(start).Milliseconds()}
	if err != nil {
		e.Error = err.Error()
	} else {
		e.Note = call.summary()
	}
	if f.log != nil {
		f.log.Write(e)
	}
	if w != nil {
		w.Landed("Sentry", call.Tokens, err)
	}
	return call, err
}