package main

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

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

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

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

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

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

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

// 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, ModelAuthorizer, error) {
	if config.Endpoint == configEndpointAPI {
		key := os.Getenv("OPENAI_API_KEY")
		if key == "" {
			return "", nil, errors.New("endpoint api needs OPENAI_API_KEY")
		}

		return modelAPIURL, ModelAPIKey(key), nil
	}

	return modelChatGPTURL, newChatGPTAuth(filepath.Join(dir, chatgptTokenFileName), chatgptIssuer), nil
}

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

// modelUserMessage returns a user input item with one text part for each text.
func modelUserMessage(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,
	}
}

// modelOutputText returns the text of the message items of output.
func modelOutputText(output []json.RawMessage) (string, error) {
	var text strings.Builder
	for _, raw := range output {
		var item ModelOutputItem
		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
}

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

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

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

// ModelClient 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.
// A call of a subagent sends the key of the agent with "-" and the name of the
// subagent, so that the calls of a subagent do not share a cache entry with the
// calls of the agent that run at the same time. 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 ModelClient struct {
	url      string
	auth     ModelAuthorizer
	cacheKey string
	usage    *ModelUsageLog
	http     *http.Client
	limit    int64
	calls    atomic.Int64
}

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

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

// CallLimit returns the configured call limit, or zero for no limit.
func (c *ModelClient) CallLimit() int64 {
	return c.limit
}

// UsagePath returns the file that records the calls of this client.
func (c *ModelClient) UsagePath() string {
	return c.usage.path
}

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

	cacheKey := c.cacheKey
	if request.Agent != "" {
		cacheKey += "-" + request.Agent
	}

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

	httpRequest, err := http.NewRequestWithContext(ctx, http.MethodPost, c.url, bytes.NewReader(body))
	if err != nil {
		return ModelResponse{}, 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 ModelResponse{}, fmt.Errorf("authorize model call: %w", err)
	}

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

	start := time.Now()
	record := ModelUsageRecord{
		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
	record.Agent = request.Agent
	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 *ModelClient) send(httpRequest *http.Request, record *ModelUsageRecord) (response ModelResponse, err error) {
	httpResponse, err := c.http.Do(httpRequest)
	if err != nil {
		return ModelResponse{}, fmt.Errorf("send model call: %w", err)
	}

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

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

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

	return modelReadStream(httpResponse.Body)
}

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

// Write appends record to the file with one write.
func (u *ModelUsageLog) Write(record ModelUsageRecord) (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
}

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

// ModelRequest is one model call. Input holds the input items in order, such as
// modelUserMessage 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. Agent names the subagent that makes the call, or is
// empty for a call of the agent itself. Node, Turn, and Agent go to the usage
// file, never to the endpoint as fields. Agent also gives the call the cache
// key of its subagent.
type ModelRequest struct {
	Model        string
	Effort       string
	Instructions string
	Tools        []ModelTool
	Input        []any
	Node         *NodeKey
	Turn         string
	Agent        string
}

// ModelResponse 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 ModelResponse struct {
	Output []json.RawMessage
	Text   string
	Usage  ModelUsage
}

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

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

// ModelUsageRecord 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. Agent names the subagent
// of the call. Quota holds the x-codex-* headers of the response, without the
// turn state.
type ModelUsageRecord struct {
	Date         string            `json:"date"`
	Seconds      float64           `json:"seconds"`
	Model        string            `json:"model"`
	Node         *[2]int           `json:"node,omitempty"`
	Turn         string            `json:"turn,omitempty"`
	Agent        string            `json:"agent,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"`
}

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

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

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

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

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

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

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

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

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

	return quota
}

// modelReadStream 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 modelReadStream(stream io.Reader) (ModelResponse, error) {
	var response ModelResponse
	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 := modelReadEvent(&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")
		}
	}
}

// modelReadEvent adds one stream event to response, and tells if it ends the
// stream.
func modelReadEvent(response *ModelResponse, data []byte) (bool, error) {
	var event ModelStreamEvent
	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 = ModelUsage{
				Input:  event.Response.Usage.InputTokens,
				Cached: event.Response.Usage.InputTokensDetails.CachedTokens,
				Output: event.Response.Usage.OutputTokens,
			}
		}

		response.Text, err = modelOutputText(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
}