// Package laya talks to the Laya decision API (https://laya.k6n.net) - a fast
// non-autoregressive "System 1" model that answers typed questions about a
// state in one forward pass. It is used here as an incident triage brain: give
// it a compact textual state of a firing spike and it returns a calibrated
// choice (restart / ignore / notify / escalate) with probabilities.
//
// The API is described by https://laya.k6n.net/openapi.json. Only the parts
// this repo uses are modelled; the wire format is deliberately loose on the
// response side (answers are a map) because the model decides which questions
// get answers.
package laya

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

// Question kinds understood by the API.
const (
	KindChoice = "choice" // pick one label from Criteria
	KindScore  = "score"  // pick a level index from Levels
	KindNoul   = "noul"   // a single 0..1 non-choice answer
)

// Question is one typed question. Use the ChoiceQuestion, ScoreQuestion or
// NoulQuestion constructors rather than building it by hand, because
// "criteria" is a JSON object for choice questions and a JSON array for score
// ones - both are rendered by MarshalJSON.
type Question struct {
	kind         string
	instructions string
	criteria     map[string]string // choice: label -> what that label means
	levels       []string          // score: level descriptions, index 0 first
}

func (q Question) MarshalJSON() ([]byte, error) {
	out := map[string]any{"type": q.kind, "instructions": q.instructions}
	switch q.kind {
	case KindChoice:
		if len(q.criteria) == 0 {
			return nil, errors.New("laya: choice question needs criteria")
		}
		out["criteria"] = q.criteria
	case KindScore:
		if len(q.levels) == 0 {
			return nil, errors.New("laya: score question needs criteria levels")
		}
		out["criteria"] = q.levels
	}
	return json.Marshal(out)
}

func (q Question) Kind() string { return q.kind }

// ChoiceQuestion asks the model to pick one of criteria. The label -> text map
// matters: the model reads the descriptions, not just the labels, so "restart"
// should say what restarting means here.
func ChoiceQuestion(instructions string, criteria map[string]string) Question {
	return Question{kind: KindChoice, instructions: instructions, criteria: criteria}
}

// ScoreQuestion asks for a level between 0 and len(levels)-1.
func ScoreQuestion(instructions string, levels []string) Question {
	return Question{kind: KindScore, instructions: instructions, levels: levels}
}

// NoulQuestion asks a single open question answered with one 0..1 value.
func NoulQuestion(instructions string) Question {
	return Question{kind: KindNoul, instructions: instructions}
}

// PredictRequest is the body of POST /v1/predict. State is a string, a map or
// a list; this repo always sends a rendered text state.
type PredictRequest struct {
	State     any                 `json:"state"`
	Questions map[string]Question `json:"questions"`
	Lang      string              `json:"lang,omitempty"`
	MaxLen    int                 `json:"max_len,omitempty"`
	BatchSize int                 `json:"batch_size,omitempty"` // batch endpoint only
}

// BatchPredictRequest is the body of POST /v1/predict/batch: several states
// evaluated against the same questions, packed into shared forward passes.
type BatchPredictRequest struct {
	States    []any               `json:"states"`
	Questions map[string]Question `json:"questions"`
	Lang      string              `json:"lang,omitempty"`
	BatchSize int                 `json:"batch_size,omitempty"`
	MaxLen    int                 `json:"max_len,omitempty"`
}

// Answer is one answered question. Only the fields matching Type are set.
type Answer struct {
	Type          string             `json:"type"`
	Choice        string             `json:"choice,omitempty"`
	Probabilities map[string]float64 `json:"probabilities,omitempty"`
	Score         float64            `json:"score,omitempty"`
	Legend        map[string]string  `json:"legend,omitempty"`
	Noul          float64            `json:"noul,omitempty"`
	Confidence    float64            `json:"confidence,omitempty"`
	// AnswerConfidence is the model's confidence in the answer it picked,
	// as opposed to Confidence, which measures how decisive the forward pass
	// was overall.
	AnswerConfidence float64 `json:"answer_confidence,omitempty"`
}

// Probability returns the model's probability for one choice label, 0 if the
// label was not offered or not answered.
func (a Answer) Probability(label string) float64 { return a.Probabilities[label] }

// Margin is how far the chosen answer leads the next best one. The label
// probabilities are diffuse - even a correct answer rarely exceeds 0.5 - so
// the lead over the runner-up is a better confidence signal than the raw
// probability.
func (a Answer) Margin() float64 {
	best, second := 0.0, 0.0
	for _, p := range a.Probabilities {
		switch {
		case p > best:
			second, best = best, p
		case p > second:
			second = p
		}
	}
	return best - second
}

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"`
}

// Health is the body of GET /healthz.
type Health struct {
	Status  string `json:"status"`
	Model   string `json:"model"`
	Device  string `json:"device"`
	Version string `json:"version"`
}

// Client is a Laya API client. BaseURL defaults to https://laya.k6n.net.
type Client struct {
	BaseURL   string
	HTTP      *http.Client
	Lang      string
	MaxLen    int
	BatchSize int // forward-pass chunking for the batch endpoint; 0 = server default
}

// NewClient returns a client for base ("https://laya.k6n.net" when empty).
func NewClient(base string) *Client {
	return &Client{BaseURL: strings.TrimRight(base, "/"), HTTP: &http.Client{Timeout: 30 * time.Second}, Lang: "en"}
}

func (c *Client) url(path string) string {
	base := c.BaseURL
	if base == "" {
		base = "https://laya.k6n.net"
	}
	return base + path
}

func (c *Client) client() *http.Client {
	if c.HTTP != nil {
		return c.HTTP
	}
	return &http.Client{Timeout: 30 * time.Second}
}

// Health reports the model behind the API, for a setup probe.
func (c *Client) Health(ctx context.Context) (Health, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.url("/healthz"), nil)
	if err != nil {
		return Health{}, err
	}
	resp, err := c.client().Do(req)
	if err != nil {
		return Health{}, unwrap(err)
	}
	defer resp.Body.Close()
	if err := checkStatus(resp); err != nil {
		return Health{}, err
	}
	var h Health
	if err := json.NewDecoder(io.LimitReader(resp.Body, 1<<16)).Decode(&h); err != nil {
		return Health{}, fmt.Errorf("laya: decode health: %w", err)
	}
	return h, nil
}

// Predict evaluates one state against the questions in a single pass.
func (c *Client) Predict(ctx context.Context, state any, questions map[string]Question) (Response, error) {
	body := PredictRequest{State: state, Questions: questions, Lang: c.Lang, MaxLen: c.MaxLen}
	return c.post(ctx, "/v1/predict", body)
}

// PredictBatch evaluates several states against the same questions. Results
// are returned in request order, one entry per state.
func (c *Client) PredictBatch(ctx context.Context, states []any, questions map[string]Question) ([]Response, error) {
	body := BatchPredictRequest{States: states, Questions: questions, Lang: c.Lang, BatchSize: c.BatchSize, MaxLen: c.MaxLen}
	var out struct {
		Results []Response `json:"results"`
	}
	if err := c.postInto(ctx, "/v1/predict/batch", body, &out); err != nil {
		return nil, err
	}
	if len(out.Results) != len(states) {
		return nil, fmt.Errorf("laya: batch returned %d results for %d states", len(out.Results), len(states))
	}
	return out.Results, nil
}

func (c *Client) post(ctx context.Context, path string, body any) (Response, error) {
	var out Response
	if err := c.postInto(ctx, path, body, &out); err != nil {
		return Response{}, err
	}
	return out, nil
}

func (c *Client) postInto(ctx context.Context, path string, body, out any) error {
	raw, err := json.Marshal(body)
	if err != nil {
		return err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.url(path), bytes.NewReader(raw))
	if err != nil {
		return errors.New("laya: invalid request")
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := c.client().Do(req)
	if err != nil {
		return unwrap(err)
	}
	defer resp.Body.Close()
	if err := checkStatus(resp); err != nil {
		return err
	}
	return json.NewDecoder(io.LimitReader(resp.Body, 1<<22)).Decode(out)
}

// unwrap drops the URL the transport puts in its error, which may carry a
// token; the returned error must be safe to log.
func unwrap(err error) error {
	var ue interface{ Unwrap() error }
	if errors.As(err, &ue) {
		return errors.New(scrub(ue.Unwrap().Error()))
	}
	return errors.New(scrub(err.Error()))
}

// scrub removes anything that looks like a credential from a URL in an error.
func scrub(msg string) string {
	for _, part := range strings.Fields(msg) {
		if i := strings.Index(part, "://"); i >= 0 {
			rest := part[i+3:]
			if at := strings.IndexByte(rest, '@'); at >= 0 {
				host := rest[at+1:]
				msg = strings.ReplaceAll(msg, part, "laya://"+host)
			}
		}
	}
	return msg
}

func checkStatus(resp *http.Response) error {
	if resp.StatusCode/100 == 2 {
		return nil
	}
	detail, _ := io.ReadAll(io.LimitReader(resp.Body, 500))
	msg := strings.TrimSpace(string(detail))
	if msg == "" {
		msg = http.StatusText(resp.StatusCode)
	}
	return fmt.Errorf("laya: HTTP %d: %s", resp.StatusCode, scrub(msg))
}