package llm

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

const (
	DefaultBaseURL = "https://api.anthropic.com"
	DefaultModel   = "claude-opus-5"
	// Non-streaming requests need to finish inside the HTTP timeout; 16k is the
	// documented ceiling for that. Phase 5 raises this once we stream.
	DefaultMaxTokens  = 16000
	apiVersion        = "2023-06-01"
	oauthBeta         = "oauth-2025-04-20"
	defaultMaxRetries = 2
)

// Client talks to POST /v1/messages.
type Client struct {
	BaseURL    string
	MaxTokens  int
	MaxRetries int
	// RetryDelay is the base of the exponential backoff between retries.
	RetryDelay time.Duration
	HTTP       *http.Client

	apiKey      string
	authToken   string
	tokenSource func(context.Context) (string, error)

	// model is behind a lock because a front end can switch it while the
	// agent is between turns. It is the one field that moves after
	// construction, so it is the one that is not public.
	mu    sync.Mutex
	model string
}

// Model reports the model requests default to.
func (c *Client) Model() string {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.model
}

// SetModel changes the model the requests that follow will ask for.
func (c *Client) SetModel(id string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.model = id
}

// Clone returns a copy that authenticates exactly as this client does and
// differs only in the model it asks for. A subagent runs on its own model, and
// must keep running on it even if the session switches the parent's — so the
// model is copied, not shared, and so is the lock that guards it.
func (c *Client) Clone(model string) *Client {
	c.mu.Lock()
	defer c.mu.Unlock()
	if model == "" {
		model = c.model
	}
	return &Client{
		BaseURL:     c.BaseURL,
		MaxTokens:   c.MaxTokens,
		MaxRetries:  c.MaxRetries,
		RetryDelay:  c.RetryDelay,
		HTTP:        c.HTTP,
		apiKey:      c.apiKey,
		authToken:   c.authToken,
		tokenSource: c.tokenSource,
		model:       model,
	}
}

// errNoCredentials is what every route reports when it has nothing to
// authenticate with. It is one value rather than three copies of a string
// because ListModels, CreateMessage and StreamMessage all reach it, and a
// caller comparing them should not have to hope they were typed alike.
var errNoCredentials = errors.New("llm: no credentials: set ANTHROPIC_API_KEY or ANTHROPIC_AUTH_TOKEN")

// Option configures a Client.
type Option func(*Client)

// WithBaseURL points the client at a different endpoint: a proxy, a gateway,
// or a local server speaking the Messages API. An empty string is ignored, so
// callers can pass an unset flag without clobbering the environment.
func WithBaseURL(u string) Option {
	return func(c *Client) {
		if u = normalizeBaseURL(u); u != "" {
			c.BaseURL = u
		}
	}
}
func WithModel(m string) Option {
	return func(c *Client) {
		if m != "" {
			c.model = m
		}
	}
}

func WithMaxTokens(n int) Option { return func(c *Client) { c.MaxTokens = n } }

// WithAPIKey sets the x-api-key credential. An empty string is ignored.
func WithAPIKey(k string) Option {
	return func(c *Client) {
		if k != "" {
			c.apiKey, c.authToken = k, ""
		}
	}
}

// WithAuthToken sets a bearer credential, which is what OAuth tokens and most
// third-party gateways expect. An empty string is ignored.
func WithAuthToken(t string) Option {
	return func(c *Client) {
		if t != "" {
			c.authToken, c.apiKey = t, ""
		}
	}
}

// WithTokenSource sets a bearer credential resolved per request, so a token
// that expires mid-session can be replaced without rebuilding the client. A
// nil function is ignored.
func WithTokenSource(f func(context.Context) (string, error)) Option {
	return func(c *Client) {
		if f != nil {
			c.tokenSource, c.apiKey, c.authToken = f, "", ""
		}
	}
}
func WithHTTPClient(h *http.Client) Option { return func(c *Client) { c.HTTP = h } }

// New builds a client, reading credentials from the environment unless an
// option overrides them. An unset ANTHROPIC_API_KEY is not fatal here: the
// caller decides, because ANTHROPIC_AUTH_TOKEN or an `ant auth login` profile
// may be in play.
func New(opts ...Option) *Client {
	c := &Client{
		BaseURL:    DefaultBaseURL,
		model:      DefaultModel,
		MaxTokens:  DefaultMaxTokens,
		MaxRetries: defaultMaxRetries,
		RetryDelay: 500 * time.Millisecond,
		HTTP:       &http.Client{Timeout: 10 * time.Minute},
		apiKey:     os.Getenv("ANTHROPIC_API_KEY"),
		authToken:  os.Getenv("ANTHROPIC_AUTH_TOKEN"),
	}
	if c.apiKey != "" {
		c.authToken = ""
	}
	if u := normalizeBaseURL(os.Getenv("ANTHROPIC_BASE_URL")); u != "" {
		c.BaseURL = u
	}
	for _, o := range opts {
		o(c)
	}
	return c
}

// normalizeBaseURL trims whitespace and trailing slashes, so a base URL pasted
// with a trailing "/" does not produce a request to "//v1/messages" — which
// some proxies 404 rather than normalize.
func normalizeBaseURL(u string) string {
	return strings.TrimRight(strings.TrimSpace(u), "/")
}

// IsDefaultEndpoint reports whether the client is talking to Anthropic.
func (c *Client) IsDefaultEndpoint() bool { return c.BaseURL == DefaultBaseURL }

// HasCredentials reports whether the client can authenticate.
func (c *Client) HasCredentials() bool {
	return c.apiKey != "" || c.authToken != "" || c.tokenSource != nil
}

// ClearCredentialsForTest drops any credential read from the environment, so a
// test exercising credential precedence is not at the mercy of the developer's
// own shell.
func (c *Client) ClearCredentialsForTest() {
	c.apiKey, c.authToken, c.tokenSource = "", "", nil
}

// APIError is a non-2xx response.
type APIError struct {
	StatusCode int
	Type       string
	Message    string
	Body       string
}

func (e *APIError) Error() string {
	if e.Message != "" {
		return fmt.Sprintf("anthropic: %d %s: %s", e.StatusCode, e.Type, e.Message)
	}
	return fmt.Sprintf("anthropic: %d: %s", e.StatusCode, e.Body)
}

// Retryable reports whether resending the identical request could succeed.
func (e *APIError) Retryable() bool {
	return e.StatusCode == http.StatusTooManyRequests ||
		e.StatusCode == http.StatusRequestTimeout ||
		e.StatusCode == http.StatusConflict ||
		e.StatusCode >= 500
}

// auth is one request's resolved credential. Resolving returns a value rather
// than mutating the client so concurrent requests cannot interleave.
type auth struct {
	apiKey string
	bearer string
}

// apply writes the credential headers. An OAuth bearer also needs the beta
// flag; an API key does not.
func (a auth) apply(h http.Header) {
	if a.apiKey != "" {
		h.Set("x-api-key", a.apiKey)
		return
	}
	h.Set("authorization", "Bearer "+a.bearer)
	h.Set("anthropic-beta", oauthBeta)
}

// resolveAuth produces the credential for one request. It runs before the
// retry loop: a credential that will not resolve is permanent, and the loop
// only short-circuits on a non-retryable *APIError, so resolving inside
// attempt would retry a dead credential three times with backoff.
func (c *Client) resolveAuth(ctx context.Context) (auth, error) {
	if c.apiKey != "" {
		return auth{apiKey: c.apiKey}, nil
	}
	if c.tokenSource != nil {
		tok, err := c.tokenSource(ctx)
		if err != nil {
			return auth{}, err
		}
		return auth{bearer: tok}, nil
	}
	return auth{bearer: c.authToken}, nil
}

// CreateMessage sends one request, filling in the client's defaults and
// retrying the errors worth retrying.
func (c *Client) CreateMessage(ctx context.Context, req Request) (*Response, error) {
	// Checked once, up front: a missing credential is permanent, and retrying
	// it three times with backoff just delays the error.
	if !c.HasCredentials() {
		return nil, errNoCredentials
	}
	a, err := c.resolveAuth(ctx)
	if err != nil {
		return nil, err
	}
	if req.Model == "" {
		req.Model = c.Model()
	}
	if req.MaxTokens == 0 {
		req.MaxTokens = c.MaxTokens
	}
	body, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("llm: encoding request: %w", err)
	}

	var lastErr error
	for attempt := 0; attempt <= c.MaxRetries; attempt++ {
		if attempt > 0 {
			if err := sleep(ctx, c.backoff(attempt, lastErr)); err != nil {
				return nil, err
			}
		}
		resp, err := c.attempt(ctx, body, a)
		if err == nil {
			return resp, nil
		}
		lastErr = err
		if ctx.Err() != nil {
			return nil, ctx.Err()
		}
		var apiErr *APIError
		if ok := asAPIError(err, &apiErr); ok && !apiErr.Retryable() {
			return nil, err
		}
	}
	return nil, lastErr
}

// newRequest builds a request to the endpoint with the version and the
// credential already on it. Every call the client makes goes through here, so
// a second route cannot end up authenticating differently from the first.
//
// The credential arrives as a value rather than being read off the client,
// because it may have to be refreshed and that has to happen once per call,
// outside the retry loop — see resolveAuth.
func (c *Client) newRequest(ctx context.Context, method, path string, body []byte, a auth) (*http.Request, error) {
	var rdr io.Reader
	if body != nil {
		rdr = bytes.NewReader(body)
	}
	req, err := http.NewRequestWithContext(ctx, method, c.BaseURL+path, rdr)
	if err != nil {
		return nil, err
	}
	req.Header.Set("content-type", "application/json")
	req.Header.Set("anthropic-version", apiVersion)
	a.apply(req.Header)
	return req, nil
}

func (c *Client) attempt(ctx context.Context, body []byte, a auth) (*Response, error) {
	httpReq, err := c.newRequest(ctx, http.MethodPost, "/v1/messages", body, a)
	if err != nil {
		return nil, err
	}

	httpResp, err := c.HTTP.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("llm: %w", err)
	}
	defer httpResp.Body.Close()

	raw, err := io.ReadAll(httpResp.Body)
	if err != nil {
		return nil, fmt.Errorf("llm: reading response: %w", err)
	}
	if httpResp.StatusCode != http.StatusOK {
		return nil, parseAPIError(httpResp.StatusCode, httpResp.Header.Get("retry-after"), raw)
	}

	var out Response
	if err := json.Unmarshal(raw, &out); err != nil {
		return nil, fmt.Errorf("llm: decoding response: %w", err)
	}
	return &out, nil
}

func asAPIError(err error, target **APIError) bool {
	if e, ok := err.(*APIError); ok {
		*target = e
		return true
	}
	return false
}

// backoff grows exponentially, but honours a retry-after hint when the server
// sent one.
func (c *Client) backoff(attempt int, lastErr error) time.Duration {
	var apiErr *APIError
	if asAPIError(lastErr, &apiErr) {
		if d, ok := retryAfter(apiErr.Message); ok {
			return d
		}
	}
	return time.Duration(math.Pow(2, float64(attempt))) * c.RetryDelay
}

func retryAfter(msg string) (time.Duration, bool) {
	const marker = "(retry-after "
	i := indexOf(msg, marker)
	if i < 0 {
		return 0, false
	}
	rest := msg[i+len(marker):]
	j := indexOf(rest, "s)")
	if j < 0 {
		return 0, false
	}
	secs, err := strconv.Atoi(rest[:j])
	if err != nil || secs < 0 {
		return 0, false
	}
	return time.Duration(secs) * time.Second, true
}

func indexOf(s, sub string) int { return bytes.Index([]byte(s), []byte(sub)) }

func sleep(ctx context.Context, d time.Duration) error {
	t := time.NewTimer(d)
	defer t.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-t.C:
		return nil
	}
}