package chatcmpl

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

// UpstreamError is returned when the upstream provider answers with a
// non-2xx status or an error payload.
type UpstreamError struct {
	StatusCode int
	Message    string
}

func (e *UpstreamError) Error() string {
	return fmt.Sprintf("upstream %d: %s", e.StatusCode, e.Message)
}

// Client talks to one Chat Completions endpoint.
type Client struct {
	BaseURL string
	APIKey  string
	HTTP    *http.Client
}

// NewClient builds a client for the given base URL and API key.
func NewClient(baseURL, apiKey string) *Client {
	return &Client{
		BaseURL: strings.TrimRight(strings.TrimSpace(baseURL), "/"),
		APIKey:  apiKey,
		HTTP:    &http.Client{Timeout: 30 * time.Minute},
	}
}

func (c *Client) endpoint() string {
	return c.BaseURL + "/chat/completions"
}

func (c *Client) do(ctx context.Context, req *Request) (*http.Response, error) {
	body, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("encode request: %w", err)
	}
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.endpoint(), bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Accept", "application/json")
	if c.APIKey != "" {
		httpReq.Header.Set("Authorization", "Bearer "+c.APIKey)
	}
	resp, err := c.HTTP.Do(httpReq)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		defer resp.Body.Close()
		msg := readErrorBody(resp)
		return nil, &UpstreamError{StatusCode: resp.StatusCode, Message: msg}
	}
	return resp, nil
}

// Create performs a non-streaming completion.
func (c *Client) Create(ctx context.Context, req *Request) (*Response, error) {
	resp, err := c.do(ctx, req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	var out Response
	if err := json.NewDecoder(io.LimitReader(resp.Body, 64<<20)).Decode(&out); err != nil {
		return nil, fmt.Errorf("decode upstream response: %w", err)
	}
	if out.Error != nil && out.Error.Message != "" {
		return nil, &UpstreamError{StatusCode: resp.StatusCode, Message: out.Error.Message}
	}
	return &out, nil
}

// CreateStream performs a streaming completion. The caller reads SSE
// events from the response body and must close it.
func (c *Client) CreateStream(ctx context.Context, req *Request) (*http.Response, error) {
	return c.do(ctx, req)
}

func readErrorBody(resp *http.Response) string {
	data, err := io.ReadAll(io.LimitReader(resp.Body, 64<<10))
	if err != nil || len(data) == 0 {
		return resp.Status
	}
	var payload struct {
		Error *Error `json:"error"`
	}
	if json.Unmarshal(data, &payload) == nil && payload.Error != nil && payload.Error.Message != "" {
		return payload.Error.Message
	}
	return strings.TrimSpace(string(data))
}