package main

// One judgement: a question answered by the node's agent with no tools (novox/hq ADR 0278 §5).
//
// A module that needs a judgement only a language model can make — the issue tracker's duplicate check — asks
// this module, which holds the agent and the licence the node is bound to, rather than taking the token itself
// (ADR 0206). The agent runs in print mode as the account agents run as on this node (ADR 0266), in an empty
// directory, with its built-in tools off, in safe mode (no tool servers, plugins, hooks or instruction files)
// and keeping no session: the session can do nothing but answer, so a question that carries text written to
// steer it can at most steer the answer. The answer is shaped by the caller's JSON schema.

import (
	"bytes"
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"

	stdio "git.novox.be/novox/mesh-sdk/go"
)

const (
	// JudgeBound is how long one judgement may take, waiting for another included: under the bus's 30 seconds
	// a tool call is answered in, so the caller hears this module's refusal rather than a timeout.
	JudgeBound = 25 * time.Second
	// MostJudgeBound is the longest bound a caller may ask with within: a long record's review takes longer than
	// one bus call allows (novox/hq issue 508), and is then collected by its ticket.
	MostJudgeBound = 120 * time.Second
	// JudgeWait is how long one call waits for a judgement asked with a bound past JudgeBound before it answers
	// the ticket: well under the bus's 30 seconds, so the call is always answered.
	JudgeWait = 20 * time.Second
	// JudgedKept is how long an answer nobody collected is kept past its judgement's bound.
	JudgedKept = 10 * time.Minute
	// MostQuestion is the largest question taken.
	MostQuestion = 64 << 10
	// MostSchema is the largest answer schema taken.
	MostSchema = 16 << 10
	// JudgeProgram is the agent's program as the module's package installs it.
	JudgeProgram = "/usr/bin/claude"
)

// JudgeModels are the model aliases a judgement may name; the first is the default.
var JudgeModels = []string{"sonnet", "haiku", "opus"}

// JudgeAsk is one question.
type JudgeAsk struct {
	Question string
	Schema   map[string]any
	Model    string
	// Within is the bound the caller asks; JudgeBound when zero. Past JudgeBound the call answers a ticket if the
	// judgement is not done within JudgeWait, and claude_code_judged collects it (novox/hq issue 508).
	Within time.Duration
}

// Judged is one answer.
type Judged struct {
	Answer  json.RawMessage `json:"answer,omitempty"`
	Model   string          `json:"model"`
	Node    string          `json:"node"`
	Seconds float64         `json:"seconds"`
	// Ticket names a judgement asked with a bound past JudgeBound; Pending says it is not done yet, and
	// claude_code_judged with the ticket collects it.
	Ticket  string `json:"ticket,omitempty"`
	Pending bool   `json:"pending,omitempty"`
}

// running is one judgement asked with a bound past JudgeBound, until it is collected.
type running struct {
	done  chan struct{}
	out   *Judged
	err   error
	until time.Time
}

// Judge runs judgements on this node, one at a time.
type Judge struct {
	P Paths
	// Licensed says whether the agent's account holds a licence now.
	Licensed func() bool
	// Exec runs the command line with the question on its input and answers what it printed.
	Exec func(ctx context.Context, argv []string, stdin []byte) ([]byte, error)
	Now  func() time.Time
	// Bound is how long one judgement may take; JudgeBound unless a test says less.
	Bound time.Duration
	// Wait is how long a call waits for a long judgement before it answers its ticket; JudgeWait unless a test
	// says less.
	Wait time.Duration
	turn chan struct{}

	mu      sync.Mutex
	tickets map[string]*running
}

// NewJudge is the judge of this node.
func NewJudge(p Paths) *Judge {
	return &Judge{P: p, Licensed: func() bool { return holdsAnyLicence(p) }, Exec: execJudgement, Now: time.Now,
		turn: make(chan struct{}, 1)}
}

// Argv is the command line of one judgement: the agent in print mode with no tools, in safe mode, keeping no
// session, as the account agents run as here — through sudo when that is not the account this runs as.
func (j *Judge) Argv(schema, model string) []string {
	argv := []string{JudgeProgram, "-p", "--output-format", "json", "--json-schema", schema, "--model", model,
		"--tools", "", "--safe-mode", "--no-session-persistence"}
	if j.P.Confined() {
		argv = append([]string{"sudo", "-n", "-H", "-u", j.P.Agent, "--"}, argv...)
	}
	return argv
}

// Judge answers one question, or refuses in words.
func (j *Judge) Judge(a JudgeAsk) (*Judged, error) {
	question := strings.TrimSpace(a.Question)
	if question == "" {
		return nil, errors.New("claude_code_judge needs a question")
	}
	if len(question) > MostQuestion {
		return nil, fmt.Errorf("the question is %d bytes; a judgement takes at most %d (64 KiB)", len(question), MostQuestion)
	}
	if a.Schema == nil {
		return nil, errors.New("claude_code_judge needs schema: the JSON schema of the answer, an object")
	}
	if t, _ := a.Schema["type"].(string); t != "object" {
		return nil, errors.New("the answer's schema is an object's: its type is \"object\"")
	}
	schema, err := json.Marshal(a.Schema)
	if err != nil || len(schema) > MostSchema {
		return nil, fmt.Errorf("the answer's schema is not JSON of at most %d bytes", MostSchema)
	}
	model := strings.TrimSpace(a.Model)
	if model == "" {
		model = JudgeModels[0]
	}
	known := false
	for _, m := range JudgeModels {
		known = known || m == model
	}
	if !known {
		return nil, fmt.Errorf("model %q is not one of %s", model, strings.Join(JudgeModels, ", "))
	}
	if !j.Licensed() {
		return nil, fmt.Errorf("the agent on %s holds no licence, so nothing can judge here: the licence manager's "+
			"bindings say which node holds which", j.P.Node)
	}

	bound := j.Bound
	if bound <= 0 {
		bound = JudgeBound
	}
	if a.Within > MostJudgeBound {
		return nil, fmt.Errorf("within is %s; a judgement takes at most %s", a.Within, MostJudgeBound)
	}
	if a.Within > bound {
		return j.start(string(schema), model, question, a.Within)
	}
	if a.Within > 0 {
		bound = a.Within
	}
	return j.run(string(schema), model, question, bound)
}

// run is one judgement within its bound, waiting for its turn included.
func (j *Judge) run(schema, model, question string, bound time.Duration) (*Judged, error) {
	start := j.Now()
	ctx, cancel := context.WithTimeout(context.Background(), bound)
	defer cancel()
	select {
	case j.turn <- struct{}{}:
		defer func() { <-j.turn }()
	case <-ctx.Done():
		return nil, fmt.Errorf("another judgement held %s's agent for all of %s; ask again", j.P.Node, bound)
	}
	out, err := j.Exec(ctx, j.Argv(schema, model), []byte(question))
	if ctx.Err() != nil {
		return nil, fmt.Errorf("the agent on %s did not answer within %s; the judgement was ended", j.P.Node, bound)
	}
	if err != nil {
		return nil, fmt.Errorf("the agent on %s could not judge: %v", j.P.Node, err)
	}
	answer, err := answerOf(out)
	if err != nil {
		return nil, fmt.Errorf("the agent on %s answered, and not in the shape asked: %v", j.P.Node, err)
	}
	return &Judged{Answer: answer, Model: model, Node: j.P.Node, Seconds: j.Now().Sub(start).Round(100 * time.Millisecond).Seconds()}, nil
}

// start runs a judgement whose bound is past what one bus call allows, and waits for it as long as a call may: its
// answer when it is done by then, its ticket otherwise (novox/hq issue 508).
func (j *Judge) start(schema, model, question string, bound time.Duration) (*Judged, error) {
	var b [12]byte
	_, _ = rand.Read(b[:])
	ticket := hex.EncodeToString(b[:])
	r := &running{done: make(chan struct{}), until: j.Now().Add(bound + JudgedKept)}
	j.mu.Lock()
	j.forget()
	if j.tickets == nil {
		j.tickets = map[string]*running{}
	}
	j.tickets[ticket] = r
	j.mu.Unlock()
	go func() {
		r.out, r.err = j.run(schema, model, question, bound)
		close(r.done)
	}()
	return j.Collect(ticket)
}

// Collect waits as long as a call may for the judgement under a ticket: its answer, why it failed, or that it is
// still pending. A ticket is answered once; an answer nobody collects is kept JudgedKept past its bound.
func (j *Judge) Collect(ticket string) (*Judged, error) {
	j.mu.Lock()
	j.forget()
	r := j.tickets[ticket]
	j.mu.Unlock()
	if r == nil {
		return nil, fmt.Errorf("no judgement on %s under ticket %q: it was collected, kept past its time, or this module "+
			"restarted since; ask the judgement again", j.P.Node, ticket)
	}
	wait := j.Wait
	if wait <= 0 {
		wait = JudgeWait
	}
	select {
	case <-r.done:
	case <-time.After(wait):
		return &Judged{Node: j.P.Node, Ticket: ticket, Pending: true}, nil
	}
	j.mu.Lock()
	delete(j.tickets, ticket)
	j.mu.Unlock()
	if r.err != nil {
		return nil, r.err
	}
	out := *r.out
	out.Ticket = ticket
	return &out, nil
}

// forget drops the answers nobody collected in time. Called with mu held.
func (j *Judge) forget() {
	now := j.Now()
	for t, r := range j.tickets {
		if now.After(r.until) {
			delete(j.tickets, t)
		}
	}
}

// answerOf reads print mode's result: the structured answer, or why there is none.
func answerOf(out []byte) (json.RawMessage, error) {
	var r struct {
		IsError    bool            `json:"is_error"`
		Subtype    string          `json:"subtype"`
		Result     string          `json:"result"`
		Structured json.RawMessage `json:"structured_output"`
	}
	if err := json.Unmarshal(bytes.TrimSpace(out), &r); err != nil {
		return nil, fmt.Errorf("its output is not print mode's result (%v)", err)
	}
	if r.IsError {
		why := r.Result
		if len(why) > 300 {
			why = why[:300]
		}
		return nil, fmt.Errorf("it ended in an error (%s): %s", r.Subtype, why)
	}
	if len(r.Structured) == 0 || string(r.Structured) == "null" {
		return nil, errors.New("it gave no structured answer")
	}
	var obj map[string]any
	if err := json.Unmarshal(r.Structured, &obj); err != nil {
		return nil, errors.New("its answer is not an object")
	}
	return r.Structured, nil
}

// execJudgement runs one judgement in an empty directory, ended when the context is.
func execJudgement(ctx context.Context, argv []string, stdin []byte) ([]byte, error) {
	dir, err := os.MkdirTemp("", "claude-code-judge-")
	if err != nil {
		return nil, err
	}
	defer os.RemoveAll(dir)
	// The agent's account must be able to stand in it.
	_ = os.Chmod(dir, 0o755)
	cmd := exec.CommandContext(ctx, argv[0], argv[1:]...)
	cmd.Dir = dir
	cmd.Stdin = bytes.NewReader(stdin)
	cmd.WaitDelay = 2 * time.Second
	var stderr bytes.Buffer
	cmd.Stderr = &stderr
	out, err := cmd.Output()
	if err != nil {
		why := strings.TrimSpace(stderr.String())
		if len(why) > 300 {
			why = why[:300]
		}
		if why == "" {
			// Print mode says an error as its result on stdout.
			if a, aerr := answerOf(out); aerr != nil && a == nil && len(out) > 0 {
				why = aerr.Error()
			}
		}
		return nil, fmt.Errorf("%v: %s", err, why)
	}
	return out, nil
}

// judgeTool is the tool.
func judgeTool(j *Judge) stdio.Tool {
	return stdio.Tool{Name: "claude_code_judge",
		Description: "One judgement by this node's agent, for a module that needs a question only a language model " +
			"can answer (novox/hq ADR 0278): give the question in plain words and the JSON schema of the answer (an " +
			"object); optionally the model (" + strings.Join(JudgeModels, ", ") + "; default " + JudgeModels[0] + "). The agent " +
			"answers as the account agents run as here, under the licence this node is bound to, with no tools, " +
			"no tool servers, no instructions and no session kept — it can only answer — within 25 seconds or " +
			"not at all, unless within asks longer (at most 120 seconds, for a long question): then, when it is not done in " +
			"20 seconds, it answers pending with a ticket, and claude_code_judged collects the answer (novox/hq issue 508). " +
			"One at a time on a node. Refused: an empty or larger than 64 KiB question, a schema that " +
			"is not an object, an unknown model, a node holding no licence. The question is never logged.",
		Input: map[string]any{
			"question": str("the question, with everything the answer needs"),
			"schema":   map[string]any{"type": "object", "description": "the JSON schema the answer follows"},
			"model":    map[string]any{"type": "string", "enum": JudgeModels, "description": "the model (default " + JudgeModels[0] + ")"},
			"within": map[string]any{"type": "number", "description": "the bound in seconds, at most 120 (default 25); " +
				"past 25 a judgement not done in 20 seconds answers pending with a ticket for claude_code_judged"},
		},
		Run: func(a map[string]any) (any, error) {
			for k := range a {
				if k != "question" && k != "schema" && k != "model" && k != "within" {
					return nil, fmt.Errorf("claude_code_judge takes question, schema, model and within; not %s", k)
				}
			}
			var ask JudgeAsk
			var ok bool
			if ask.Question, ok = a["question"].(string); !ok {
				return nil, errors.New("question is text")
			}
			if v, given := a["schema"]; given {
				if ask.Schema, ok = v.(map[string]any); !ok {
					return nil, errors.New("schema is a JSON schema: an object")
				}
			}
			if v, given := a["model"]; given && v != nil {
				if ask.Model, ok = v.(string); !ok {
					return nil, errors.New("model is text")
				}
			}
			if v, given := a["within"]; given && v != nil {
				secs, ok := v.(float64)
				if !ok || secs <= 0 {
					return nil, errors.New("within is a number of seconds")
				}
				ask.Within = time.Duration(secs * float64(time.Second))
			}
			return j.Judge(ask)
		}}
}

// judgedTool collects a judgement asked with a bound past one bus call's (novox/hq issue 508).
func judgedTool(j *Judge) stdio.Tool {
	return stdio.Tool{Name: "claude_code_judged",
		Description: "The answer of a judgement claude_code_judge answered pending with a ticket (novox/hq issue 508): " +
			"waits up to 20 seconds, then answers the judgement's answer, why it failed, or pending again. A ticket is " +
			"answered once; one nobody collects is dropped ten minutes after its bound, and none survives a restart of " +
			"this module.",
		Input: map[string]any{"ticket": str("the ticket claude_code_judge answered")},
		Run: func(a map[string]any) (any, error) {
			for k := range a {
				if k != "ticket" {
					return nil, fmt.Errorf("claude_code_judged takes ticket; not %s", k)
				}
			}
			ticket, ok := a["ticket"].(string)
			if !ok || ticket == "" {
				return nil, errors.New("claude_code_judged needs ticket, as claude_code_judge answered it")
			}
			return j.Collect(ticket)
		}}
}