package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

const (
	// chatgptURL is the Responses endpoint of the ChatGPT backend.
	chatgptURL = "https://chatgpt.com/backend-api/codex/responses"

	// apiURL is the Responses endpoint of the OpenAI API.
	apiURL = "https://api.openai.com/v1/responses"

	// quotaPrefix starts the names of the quota headers of the ChatGPT
	// backend, in canonical form.
	quotaPrefix = "X-Codex-"

	// turnStateHeader is an x-codex-* header that holds an opaque state of
	// the backend, not quota.
	turnStateHeader = "X-Codex-Turn-State"

	// errorBytes is the most bytes of an error response that go into the
	// error.
	errorBytes = 4096
)

// errCallLimit tells that a client sent its limit of calls.
var errCallLimit = errors.New("call limit reached")

// authorizer adds the credentials of a model call to the headers of its
// request.
type authorizer interface {
	Authorize(ctx context.Context, header http.Header) error
}

// apiKey authorizes calls to the OpenAI API with a key.
type apiKey string

// Client sends model calls to a Responses endpoint for one agent. The ChatGPT
// backend and the OpenAI API take the same request body and differ only in the
// URL and the credentials. Every call sends the fixed cache key of the agent,
// so the turns and the compactions of the agent share one prompt cache. The
// ChatGPT backend caches only with the key in the session_id header, so a call
// to it sends the key there and in the body. Every call is recorded in the
// usage file.
//
// The request never holds max_output_tokens or metadata, because the ChatGPT
// backend refuses them. It sets store to false, so the client sends the whole
// input on each call, and it asks for the encrypted reasoning, so a later call
// can send the reasoning items back.
type Client struct {
	url      string
	auth     authorizer
	cacheKey string
	usage    *usageLog
	http     *http.Client
	limit    int64
	calls    atomic.Int64
}

// Request is one model call. Input holds the input items in order, such as
// userMessage values or output items of an earlier response. Node names the
// node that a compaction builds, or is nil. Turn names the request of a turn:
// first or next. Node and Turn go only to the usage file, never to the
// endpoint.
type Request struct {
	Model        string
	Effort       string
	Instructions string
	Tools        []Tool
	Input        []any
	Node         *nodeKey
	Turn         string
}

// Tool is the definition of a function tool.
type Tool struct {
	Type        string `json:"type"`
	Name        string `json:"name"`
	Description string `json:"description"`
	Parameters  any    `json:"parameters"`
	Strict      bool   `json:"strict"`
}

// Response is the result of a model call. Output holds the output items as the
// endpoint sent them. Text is the text of the message items.
type Response struct {
	Output []json.RawMessage
	Text   string
	Usage  Usage
}

// Usage gives the token counts of a call. Cached is the part of the input
// tokens that came from the prompt cache.
type Usage struct {
	Input  int
	Cached int
	Output int
}

// requestBody is the JSON body of a call.
type requestBody struct {
	Model          string    `json:"model"`
	Instructions   string    `json:"instructions"`
	Input          []any     `json:"input"`
	Tools          []Tool    `json:"tools"`
	Reasoning      reasoning `json:"reasoning"`
	Store          bool      `json:"store"`
	Stream         bool      `json:"stream"`
	Include        []string  `json:"include"`
	PromptCacheKey string    `json:"prompt_cache_key"`
}

type reasoning struct {
	Effort string `json:"effort"`
}

// streamEvent holds the fields of the stream events that the client reads.
type streamEvent struct {
	Type     string          `json:"type"`
	Item     json.RawMessage `json:"item"`
	Response *eventResponse  `json:"response"`
	Code     string          `json:"code"`
	Message  string          `json:"message"`
}

type eventResponse struct {
	Output            []json.RawMessage `json:"output"`
	Usage             *wireUsage        `json:"usage"`
	Error             *wireError        `json:"error"`
	IncompleteDetails *struct {
		Reason string `json:"reason"`
	} `json:"incomplete_details"`
}

type wireUsage struct {
	InputTokens        int `json:"input_tokens"`
	OutputTokens       int `json:"output_tokens"`
	InputTokensDetails struct {
		CachedTokens int `json:"cached_tokens"`
	} `json:"input_tokens_details"`
}

type wireError struct {
	Code    string `json:"code"`
	Message string `json:"message"`
}

// outputItem holds the fields of an output item that give its text.
type outputItem struct {
	Type    string `json:"type"`
	Content []struct {
		Type string `json:"type"`
		Text string `json:"text"`
	} `json:"content"`
}

// modelEndpoint returns the URL and the authorizer of the endpoint in config.
// The ChatGPT backend uses the token file in dir. The OpenAI API uses the key in
// OPENAI_API_KEY, and there is no fallback from one to the other. All clients
// share the authorizer, so one refresh runs at a time.
func modelEndpoint(config Config, dir string) (string, authorizer, error) {
	if config.Endpoint == endpointAPI {
		key := os.Getenv("OPENAI_API_KEY")
		if key == "" {
			return "", nil, errors.New("endpoint api needs OPENAI_API_KEY")
		}

		return apiURL, apiKey(key), nil
	}

	return chatgptURL, newChatGPTAuth(filepath.Join(dir, tokenFileName), authIssuer), nil
}

// newClient returns a client for the Responses endpoint at url. It records
// each call in the usage file at usagePath.
func newClient(url string, auth authorizer, cacheKey, usagePath string) *Client {
	return &Client{
		url:      url,
		auth:     auth,
		cacheKey: cacheKey,
		usage: &usageLog{
			path: usagePath,
		},
		http: &http.Client{},
	}
}

// LimitCalls makes the client refuse each call after the first limit calls.
func (c *Client) LimitCalls(limit int64) {
	c.limit = limit
}

// LimitReached tells if the client refuses further calls.
func (c *Client) LimitReached() bool {
	return c.limit > 0 && c.calls.Load() >= c.limit
}

// Create sends the call and reads its event stream to the end. Over the call
// limit it sends nothing and returns errCallLimit.
func (c *Client) Create(ctx context.Context, request Request) (Response, error) {
	if c.calls.Add(1) > c.limit && c.limit > 0 {
		return Response{}, errCallLimit
	}

	body, err := json.Marshal(requestBody{
		Model:        request.Model,
		Instructions: request.Instructions,
		Input:        request.Input,
		Tools:        request.Tools,
		Reasoning: reasoning{
			Effort: request.Effort,
		},
		Store:          false,
		Stream:         true,
		Include:        []string{"reasoning.encrypted_content"},
		PromptCacheKey: c.cacheKey,
	})
	if err != nil {
		return Response{}, fmt.Errorf("encode request: %w", err)
	}

	httpRequest, err := http.NewRequestWithContext(ctx, http.MethodPost, c.url, bytes.NewReader(body))
	if err != nil {
		return Response{}, fmt.Errorf("create request: %w", err)
	}

	httpRequest.Header.Set("Content-Type", "application/json")
	httpRequest.Header.Set("Accept", "text/event-stream")
	err = c.auth.Authorize(ctx, httpRequest.Header)
	if err != nil {
		return Response{}, fmt.Errorf("authorize model call: %w", err)
	}

	_, chatgpt := c.auth.(*chatgptAuth)
	if chatgpt {
		httpRequest.Header.Set("session_id", c.cacheKey)
	}

	start := time.Now()
	record := usageRecord{
		Date:  start.Format(time.RFC3339Nano),
		Model: request.Model,
	}

	if request.Node != nil {
		record.Node = &[2]int{request.Node.level, request.Node.index}
	}

	record.Turn = request.Turn
	response, err := c.send(httpRequest, &record)
	if err != nil {
		record.Error = err.Error()
	}

	record.Seconds = time.Since(start).Seconds()
	record.InputTokens = response.Usage.Input
	record.CachedTokens = response.Usage.Cached
	record.OutputTokens = response.Usage.Output
	record.TextBytes = len(response.Text)

	return response, errors.Join(err, c.usage.Write(record))
}

// send makes the HTTP call and reads its stream. It puts the status and the
// quota headers into record.
func (c *Client) send(httpRequest *http.Request, record *usageRecord) (response Response, err error) {
	httpResponse, err := c.http.Do(httpRequest)
	if err != nil {
		return Response{}, fmt.Errorf("send model call: %w", err)
	}

	defer func() {
		err = errors.Join(err, httpResponse.Body.Close())
	}()

	record.Status = httpResponse.StatusCode
	record.Quota = quotaHeaders(httpResponse.Header)
	if httpResponse.StatusCode != http.StatusOK {
		text, readErr := io.ReadAll(io.LimitReader(httpResponse.Body, errorBytes))
		err = fmt.Errorf("model call: HTTP %d: %s", httpResponse.StatusCode, strings.TrimSpace(string(text)))

		return Response{}, errors.Join(err, readErr)
	}

	return readStream(httpResponse.Body)
}

// usageLog appends one JSON line for each model call to usage.jsonl. Calls can
// run at once, so a mutex serializes the writes.
type usageLog struct {
	mutex sync.Mutex
	path  string
}

// usageRecord is one line of usage.jsonl. Date is the start of the call, and
// Seconds its length. Node is the node [l, i] of a compaction. Status is the
// HTTP status, or 0 when no response came. TextBytes is the size of the text of
// the reply. Turn names the request of a turn. Quota holds the x-codex-* headers of the response, without the
// turn state.
type usageRecord struct {
	Date         string            `json:"date"`
	Seconds      float64           `json:"seconds"`
	Model        string            `json:"model"`
	Node         *[2]int           `json:"node,omitempty"`
	Turn         string            `json:"turn,omitempty"`
	Status       int               `json:"status"`
	InputTokens  int               `json:"input_tokens"`
	CachedTokens int               `json:"cached_tokens"`
	OutputTokens int               `json:"output_tokens"`
	TextBytes    int               `json:"text_bytes"`
	Quota        map[string]string `json:"quota,omitempty"`
	Error        string            `json:"error,omitempty"`
}

// Write appends record to the file with one write.
func (u *usageLog) Write(record usageRecord) (err error) {
	line, err := json.Marshal(record)
	if err != nil {
		return fmt.Errorf("encode usage: %w", err)
	}

	u.mutex.Lock()
	defer u.mutex.Unlock()

	file, err := os.OpenFile(u.path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600)
	if err != nil {
		return fmt.Errorf("open usage file: %w", err)
	}

	defer func() {
		err = errors.Join(err, file.Close())
	}()

	_, err = file.Write(append(line, '\n'))
	if err != nil {
		return fmt.Errorf("write usage: %w", err)
	}

	return nil
}

// Authorize sets the key as the bearer token.
func (k apiKey) Authorize(ctx context.Context, header http.Header) error {
	header.Set("Authorization", "Bearer "+string(k))
	return ctx.Err()
}

// userMessage returns a user input item with one text part for each text.
func userMessage(texts ...string) any {
	type part struct {
		Type string `json:"type"`
		Text string `json:"text"`
	}

	content := make([]part, len(texts))
	for n, text := range texts {
		content[n] = part{
			Type: "input_text",
			Text: text,
		}
	}

	return struct {
		Role    string `json:"role"`
		Content []part `json:"content"`
	}{
		Role:    "user",
		Content: content,
	}
}

// quotaHeaders returns the x-codex-* headers, with lowercase names, without the
// turn state.
func quotaHeaders(header http.Header) map[string]string {
	var quota map[string]string
	for name, values := range header {
		if !strings.HasPrefix(name, quotaPrefix) || name == turnStateHeader || len(values) == 0 {
			continue
		}

		if quota == nil {
			quota = make(map[string]string)
		}

		quota[strings.ToLower(name)] = values[0]
	}

	return quota
}

// readStream reads the server-sent events of a call until the response
// completes. The output items come from the output_item.done events, or from
// the completed response when the stream sent none.
func readStream(stream io.Reader) (Response, error) {
	var response Response
	reader := bufio.NewReader(stream)
	var data bytes.Buffer
	for {
		line, err := reader.ReadBytes('\n')
		if err != nil && !errors.Is(err, io.EOF) {
			return response, fmt.Errorf("read event stream: %w", err)
		}

		line = bytes.TrimRight(line, "\r\n")
		if len(line) > 0 {
			value, found := bytes.CutPrefix(line, []byte("data:"))
			if found {
				data.Write(bytes.TrimPrefix(value, []byte(" ")))
				data.WriteByte('\n')
			}
		} else if data.Len() > 0 {
			done, eventErr := readEvent(&response, data.Bytes())
			data.Reset()
			if eventErr != nil || done {
				return response, eventErr
			}
		}

		if errors.Is(err, io.EOF) {
			return response, errors.New("event stream ended before the response completed")
		}
	}
}

// readEvent adds one stream event to response, and tells if it ends the
// stream.
func readEvent(response *Response, data []byte) (bool, error) {
	var event streamEvent
	err := json.Unmarshal(data, &event)
	if err != nil {
		return true, fmt.Errorf("decode stream event: %w", err)
	}

	switch event.Type {
	case "response.output_item.done":
		response.Output = append(response.Output, event.Item)
	case "response.completed":
		if event.Response == nil {
			return true, errors.New("completed event without a response")
		}

		if len(response.Output) == 0 {
			response.Output = event.Response.Output
		}

		if event.Response.Usage != nil {
			response.Usage = Usage{
				Input:  event.Response.Usage.InputTokens,
				Cached: event.Response.Usage.InputTokensDetails.CachedTokens,
				Output: event.Response.Usage.OutputTokens,
			}
		}

		response.Text, err = outputText(response.Output)

		return true, err
	case "response.incomplete":
		reason := "unknown"
		if event.Response != nil && event.Response.IncompleteDetails != nil {
			reason = event.Response.IncompleteDetails.Reason
		}

		return true, fmt.Errorf("response incomplete: %s", reason)
	case "response.failed":
		message := "no details"
		if event.Response != nil && event.Response.Error != nil {
			message = event.Response.Error.Code + ": " + event.Response.Error.Message
		}

		return true, fmt.Errorf("response failed: %s", message)
	case "error":
		return true, fmt.Errorf("stream error: %s: %s", event.Code, event.Message)
	}

	return false, nil
}

// outputText returns the text of the message items of output.
func outputText(output []json.RawMessage) (string, error) {
	var text strings.Builder
	for _, raw := range output {
		var item outputItem
		err := json.Unmarshal(raw, &item)
		if err != nil {
			return "", fmt.Errorf("decode output item: %w", err)
		}

		if item.Type != "message" {
			continue
		}

		for _, part := range item.Content {
			if part.Type == "output_text" {
				text.WriteString(part.Text)
			}
		}
	}

	return text.String(), nil
}