// Package inference is the typed Go client for the shared, stateless MEGAMIND inference plane. It is how
// the FIPS API binary reaches embeddings, reranking, generation, and faithfulness checking WITHOUT
// hosting any ML in process: the heavy models run out of process on the inference plane (Apple Silicon,
// Metal), and this client speaks plain HTTP/JSON to them with timeouts, bounded retries, and a circuit
// breaker so a sick plane degrades the product gracefully instead of hanging it.
//
// The client sends no client identifier to the plane on the inference paths (it sends text in, gets
// vectors or text out), and the one teaching path is gated by the scrub boundary, so MEGAMIND learns
// only the craft of the work and persists nothing about any tenant.
package inference

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

	"github.com/missionworkforce/missionmatch/internal/platform/scrub"
)

// Config configures the client. The model names and dimension come from server side config so a model
// swap (for example all-minilm to bge-small-en-v1.5, both 384) is a config change, not a code change.
type Config struct {
	BaseURL         string        // MEGAMIND gateway base, e.g. http://100.94.8.94:8090
	EmbedModel      string        // 384 dim embedder (all-minilm or bge-small-en-v1.5)
	RerankModel     string        // cross encoder (bge-reranker-v2-m3) or cosine fallback
	GenerateModel   string        // local instruct model (e.g. llama3.1:8b)
	Dimension       int           // expected embedding width; a mismatch is a hard error (column is vector(384))
	Timeout         time.Duration // per request
	MaxRetries      int           // additional attempts on a transient failure
	BreakerTrips    int           // consecutive failures before the breaker opens
	BreakerCool     time.Duration // how long the breaker stays open
	TeachingEnabled bool          // when false, Teach validates the boundary but does not transmit
}

// Client talks to the MEGAMIND plane. It is safe for concurrent use.
type Client struct {
	cfg  Config
	http *http.Client
	cb   *breaker
}

// New builds a client. BaseURL is required; sensible defaults fill the rest.
func New(cfg Config) (*Client, error) {
	if cfg.BaseURL == "" {
		return nil, errors.New("inference: BaseURL is required")
	}
	if cfg.Dimension <= 0 {
		cfg.Dimension = 384
	}
	if cfg.Timeout <= 0 {
		cfg.Timeout = 20 * time.Second
	}
	if cfg.BreakerTrips <= 0 {
		cfg.BreakerTrips = 5
	}
	if cfg.BreakerCool <= 0 {
		cfg.BreakerCool = 30 * time.Second
	}
	return &Client{
		cfg:  cfg,
		http: &http.Client{Timeout: cfg.Timeout},
		cb:   newBreaker(cfg.BreakerTrips, cfg.BreakerCool),
	}, nil
}

// Health pings the plane.
func (c *Client) Health(ctx context.Context) error {
	return c.do(ctx, http.MethodGet, "/healthz", nil, nil)
}

// Embed returns one 384 dim vector per input. It hard fails if the plane returns the wrong width, since a
// width mismatch would silently corrupt the pgvector(384) columns.
func (c *Client) Embed(ctx context.Context, inputs []string) (*EmbedResponse, error) {
	if len(inputs) == 0 {
		return &EmbedResponse{Model: c.cfg.EmbedModel, Dimension: c.cfg.Dimension, Embeddings: nil}, nil
	}
	var out EmbedResponse
	if err := c.do(ctx, http.MethodPost, "/v1/embed", EmbedRequest{Model: c.cfg.EmbedModel, Inputs: inputs}, &out); err != nil {
		return nil, err
	}
	if len(out.Embeddings) != len(inputs) {
		return nil, fmt.Errorf("inference: embed returned %d vectors for %d inputs", len(out.Embeddings), len(inputs))
	}
	for i, v := range out.Embeddings {
		if len(v) != c.cfg.Dimension {
			return nil, fmt.Errorf("inference: embedding %d has width %d, expected %d (the column is vector(%d))", i, len(v), c.cfg.Dimension, c.cfg.Dimension)
		}
	}
	return &out, nil
}

// Rerank scores documents against a query, highest first.
func (c *Client) Rerank(ctx context.Context, query string, documents []string, topN int) (*RerankResponse, error) {
	var out RerankResponse
	if err := c.do(ctx, http.MethodPost, "/v1/rerank", RerankRequest{Model: c.cfg.RerankModel, Query: query, Documents: documents, TopN: topN}, &out); err != nil {
		return nil, err
	}
	return &out, nil
}

// Generate runs one stateless generation.
func (c *Client) Generate(ctx context.Context, req GenerateRequest) (*GenerateResponse, error) {
	if req.Model == "" {
		req.Model = c.cfg.GenerateModel
	}
	var out GenerateResponse
	if err := c.do(ctx, http.MethodPost, "/v1/generate", req, &out); err != nil {
		return nil, err
	}
	return &out, nil
}

// CheckFaithfulness asks the plane whether an answer is supported by its evidence (the self correction gate).
func (c *Client) CheckFaithfulness(ctx context.Context, answer string, evidence []string) (*FaithfulnessResponse, error) {
	var out FaithfulnessResponse
	if err := c.do(ctx, http.MethodPost, "/v1/faithfulness", FaithfulnessRequest{Answer: answer, Evidence: evidence}, &out); err != nil {
		return nil, err
	}
	return &out, nil
}

// Teach sends a contents scrubbed craft signal to the plane's teaching channel. The boundary is enforced
// here (scrub.Marshal refuses any client fact) AND again on the gateway. When TeachingEnabled is false
// the signal is validated but not transmitted, which is the default posture until multiple live clients
// exist (the federated loop is scaffolded, not active).
func (c *Client) Teach(ctx context.Context, signal scrub.TeachingSignal) error {
	body, err := scrub.Marshal(signal) // refuses to emit a payload carrying a client fact
	if err != nil {
		return err
	}
	if !c.cfg.TeachingEnabled {
		return nil // boundary checked, transmission disabled by posture
	}
	return c.doRaw(ctx, http.MethodPost, "/v1/teach", body, nil)
}

// do marshals req (if any), runs the request through the breaker with bounded retries, and decodes resp.
func (c *Client) do(ctx context.Context, method, path string, req any, resp any) error {
	var body []byte
	if req != nil {
		var err error
		if body, err = json.Marshal(req); err != nil {
			return err
		}
	}
	return c.doRaw(ctx, method, path, body, resp)
}

// ErrPlaneUnavailable marks that the inference plane could not be reached after retries (a network error,
// a timeout, or a 5xx) even though the circuit breaker has not opened yet. The caller maps it to a 503, so a
// plane outage degrades to a typed error, never a generic 4xx and never a hang.
var ErrPlaneUnavailable = errors.New("inference: plane unavailable")

func (c *Client) doRaw(ctx context.Context, method, path string, body []byte, resp any) error {
	if !c.cb.allow() {
		return ErrCircuitOpen
	}
	var lastErr error
	for attempt := 0; attempt <= c.cfg.MaxRetries; attempt++ {
		if attempt > 0 {
			// exponential backoff, cancellable
			select {
			case <-ctx.Done():
				return ctx.Err()
			case <-time.After(time.Duration(attempt) * 150 * time.Millisecond):
			}
		}
		err := c.attempt(ctx, method, path, body, resp)
		if err == nil {
			c.cb.record(true)
			return nil
		}
		lastErr = err
		if !transient(err) {
			c.cb.record(false)
			return err
		}
	}
	c.cb.record(false)
	// All retries exhausted on transient failures without tripping the breaker yet: a plane outage. Wrap the
	// typed sentinel so the caller returns a 503, never a generic 4xx and never a hang.
	return fmt.Errorf("%w: %s %s failed after %d attempts: %v", ErrPlaneUnavailable, method, path, c.cfg.MaxRetries+1, lastErr)
}

func (c *Client) attempt(ctx context.Context, method, path string, body []byte, resp any) error {
	var r io.Reader
	if body != nil {
		r = bytes.NewReader(body)
	}
	httpReq, err := http.NewRequestWithContext(ctx, method, c.cfg.BaseURL+path, r)
	if err != nil {
		return err
	}
	if body != nil {
		httpReq.Header.Set("Content-Type", "application/json")
	}
	res, err := c.http.Do(httpReq)
	if err != nil {
		return &transientError{err}
	}
	defer func() { _, _ = io.Copy(io.Discard, res.Body); _ = res.Body.Close() }()
	if res.StatusCode >= 500 {
		return &transientError{fmt.Errorf("inference: plane returned %d", res.StatusCode)}
	}
	if res.StatusCode >= 400 {
		snippet, _ := io.ReadAll(io.LimitReader(res.Body, 512))
		return fmt.Errorf("inference: %s %s: status %d: %s", method, path, res.StatusCode, bytes.TrimSpace(snippet))
	}
	if resp != nil {
		return json.NewDecoder(res.Body).Decode(resp)
	}
	return nil
}

// transientError marks a failure as retryable (network or 5xx).
type transientError struct{ err error }

func (e *transientError) Error() string { return e.err.Error() }
func (e *transientError) Unwrap() error { return e.err }

func transient(err error) bool {
	var t *transientError
	return errors.As(err, &t)
}