// Package jev is a minimal client for TypeSafe's System One API (model Jev): POST /v1/systemone with a state and a map
// of typed questions (noul | choice | score), answered in parallel with calibrated probabilities. API reference:
// https://docs.typesafe.ai/api (checked 2026-10-08). Used only in shadow mode here: answers are recorded next to the
// LLM diagnosis and never change what the agent does.
package jev

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

var ErrNoKey = errors.New("typesafe api key is empty")

type Question struct {
	Type         string `json:"type"`               // noul | choice | score
	Instructions string `json:"instructions"`       // the full question (the map key is not sent to the model)
	Criteria     any    `json:"criteria,omitempty"` // noul: {true:..., false:...}; choice: map[option]description; score: []level
}

type Answer struct {
	Type          string             `json:"type"`
	Noul          *float64           `json:"noul,omitempty"`
	Choice        string             `json:"choice,omitempty"`
	Score         *float64           `json:"score,omitempty"`
	Probabilities map[string]float64 `json:"probabilities,omitempty"`
	Confidence    *float64           `json:"confidence,omitempty"`
}

type Usage struct {
	InputTokens  int `json:"input_tokens"`
	OutputTokens int `json:"output_tokens"`
}

type Response struct {
	Model   string            `json:"model"`
	Answers map[string]Answer `json:"answers"`
	Usage   Usage             `json:"usage"`
}

type Client struct {
	base  string
	model string
	key   func() (string, error)
	http  *http.Client
}

// New builds a client. base defaults to https://api.typesafe.ai; model should be pinned (jev-1.13.0) so recorded
// answers stay comparable across campaigns (models page: aliases move on release).
func New(base, model string, key func() (string, error)) *Client {
	if base == "" {
		base = "https://api.typesafe.ai"
	}
	if model == "" {
		model = "jev-1.13.0"
	}
	return &Client{base: strings.TrimRight(base, "/"), model: model, key: key, http: &http.Client{Timeout: 10 * time.Second}}
}

func (c *Client) Model() string { return c.model }

// KeyPresent reports whether a non-empty key can be read (for /selfcheck).
func (c *Client) KeyPresent() bool {
	k, err := c.key()
	return err == nil && strings.TrimSpace(k) != ""
}

// Ask sends one request. 429/529 are retried once after the server's retry-after (or 1s); other errors return as is.
func (c *Client) Ask(ctx context.Context, state any, questions map[string]Question) (*Response, error) {
	key, err := c.key()
	if err != nil || strings.TrimSpace(key) == "" {
		return nil, ErrNoKey
	}
	body, _ := json.Marshal(map[string]any{"model": c.model, "state": state, "questions": questions})
	var last error
	for attempt := 0; attempt < 2; attempt++ {
		req, _ := http.NewRequestWithContext(ctx, "POST", c.base+"/v1/systemone", bytes.NewReader(body))
		req.Header.Set("Authorization", "Bearer "+strings.TrimSpace(key))
		req.Header.Set("Content-Type", "application/json")
		resp, err := c.http.Do(req)
		if err != nil {
			return nil, err
		}
		rb, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
		resp.Body.Close()
		if resp.StatusCode == 429 || resp.StatusCode == 529 {
			last = fmt.Errorf("typesafe status %d: %s", resp.StatusCode, truncate(string(rb), 200))
			wait := time.Second
			if ra := resp.Header.Get("retry-after"); ra != "" {
				if d, perr := time.ParseDuration(ra + "s"); perr == nil && d < 5*time.Second {
					wait = d
				}
			}
			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			case <-time.After(wait):
			}
			continue
		}
		if resp.StatusCode/100 != 2 {
			return nil, fmt.Errorf("typesafe status %d: %s", resp.StatusCode, truncate(string(rb), 300))
		}
		var out Response
		if err := json.Unmarshal(rb, &out); err != nil {
			return nil, fmt.Errorf("typesafe decode: %w", err)
		}
		return &out, nil
	}
	return nil, last
}

func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "..."
}