package codex

import (
	"bytes"
	"context"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/hex"
	"encoding/json"
	"errors"
	"io"
	"net"
	"net/http"
	"os"
	"strings"
	"time"
)

const (
	maxRequest = 32 << 10
	maxOutput  = 512 << 10
	timeout    = 90 * time.Second
)

type task struct {
	Prompt string `json:"prompt"`
}

func writeJSON(w http.ResponseWriter, status int, value any) {
	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Cache-Control", "no-store")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(value)
}

func fail(w http.ResponseWriter, status int, code string) {
	writeJSON(w, status, map[string]string{"error": code})
}

func readTask(w http.ResponseWriter, r *http.Request, limit int64) (task, bool) {
	var input task
	decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, limit))
	decoder.DisallowUnknownFields()
	err := decoder.Decode(&input)
	if err == nil {
		err = decoder.Decode(new(any))
		if err == io.EOF && strings.TrimSpace(input.Prompt) != "" {
			return input, true
		}
	}
	var tooLarge *http.MaxBytesError
	if errors.As(err, &tooLarge) {
		fail(w, 413, "request_too_large")
	} else {
		fail(w, 400, "invalid_task")
	}
	return task{}, false
}

// Register enables the public endpoint only when both configuration files are supplied.
// Each credential digest belongs to exactly one project.
func Register(mux *http.ServeMux, projects map[string]string, socket, grantsFile string) error {
	if socket == "" && grantsFile == "" {
		return nil
	}
	if socket == "" || grantsFile == "" {
		return errors.New("Codex requires both a socket and a grants file")
	}
	transport := &http.Transport{
		DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) {
			return (&net.Dialer{}).DialContext(ctx, "unix", socket)
		},
		DisableKeepAlives: true,
	}
	client := &http.Client{
		Transport:     transport,
		Timeout:       timeout + 2*time.Second,
		CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse },
	}
	mux.HandleFunc("POST /projects/{project}/codex/exec", func(w http.ResponseWriter, r *http.Request) {
		scheme, token, _ := strings.Cut(r.Header.Get("Authorization"), " ")
		key, err := hex.DecodeString(token)
		if !strings.EqualFold(scheme, "Bearer") || err != nil || len(key) != 32 {
			w.Header().Set("WWW-Authenticate", "Bearer")
			fail(w, 401, "unauthorized")
			return
		}
		data, err := os.ReadFile(grantsFile)
		var grants map[string]string
		if err != nil || json.Unmarshal(data, &grants) != nil {
			fail(w, 503, "codex_unavailable")
			return
		}
		hash := sha256.Sum256(key)
		var scope string
		for digest, project := range grants {
			value, err := hex.DecodeString(digest)
			if err == nil && subtle.ConstantTimeCompare(hash[:], value) == 1 {
				scope = project
			}
		}
		if scope == "" {
			w.Header().Set("WWW-Authenticate", "Bearer")
			fail(w, 401, "unauthorized")
			return
		}
		project := r.PathValue("project")
		if _, exists := projects[project]; !exists || scope != project {
			fail(w, 403, "forbidden")
			return
		}
		input, ok := readTask(w, r, maxRequest)
		if !ok {
			return
		}
		body, _ := json.Marshal(input)
		req, err := http.NewRequestWithContext(r.Context(), http.MethodPost, "http://codex/exec", bytes.NewReader(body))
		if err != nil {
			fail(w, 502, "codex_unavailable")
			return
		}
		resp, err := client.Do(req)
		if err != nil {
			if errors.Is(err, context.DeadlineExceeded) {
				fail(w, 504, "codex_timeout")
			} else {
				fail(w, 502, "codex_unavailable")
			}
			return
		}
		defer resp.Body.Close()
		data, err = io.ReadAll(io.LimitReader(resp.Body, maxOutput+1))
		if err != nil || len(data) > maxOutput || !json.Valid(data) {
			fail(w, 502, "invalid_codex_response")
			return
		}
		if retry := resp.Header.Get("Retry-After"); retry != "" {
			w.Header().Set("Retry-After", retry)
		}
		writeJSON(w, resp.StatusCode, json.RawMessage(data))
	})
	return nil
}