package codex

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"
)

const executable = "/opt/codex/node_modules/.bin/codex"

var errOutputLimit = errors.New("Codex output limit exceeded")

type outputBuffer struct {
	buffer   bytes.Buffer
	exceeded bool
	limit    int
}

func (b *outputBuffer) Write(p []byte) (int, error) {
	limit := b.limit
	if limit == 0 {
		limit = maxOutput
	}
	if len(p) > limit-b.buffer.Len() {
		b.exceeded = true
		return 0, errOutputLimit
	}
	return b.buffer.Write(p)
}

func ServeWorker(socket string) error {
	syscall.Umask(0077)
	if err := os.Remove(socket); err != nil && !errors.Is(err, os.ErrNotExist) {
		return err
	}
	listener, err := net.Listen("unix", socket)
	if err != nil {
		return err
	}
	defer listener.Close()
	if err := os.Chmod(socket, 0600); err != nil {
		return err
	}
	slot := make(chan struct{}, 1)
	mux := http.NewServeMux()
	mux.HandleFunc("POST /exec", func(w http.ResponseWriter, r *http.Request) {
		select {
		case slot <- struct{}{}:
			defer func() { <-slot }()
		default:
			w.Header().Set("Retry-After", "5")
			fail(w, 429, "codex_busy")
			return
		}
		input, ok := readTask(w, r, 6*maxRequest)
		if ok {
			execute(w, r, input.Prompt)
		}
	})
	mux.HandleFunc("GET /canvas/models", func(w http.ResponseWriter, r *http.Request) {
		models, err := workerCanvasModels()
		if err != nil {
			fail(w, 503, "codex_unavailable")
			return
		}
		writeJSON(w, 200, map[string]any{"models": models})
	})
	mux.HandleFunc("POST /canvas/exec", func(w http.ResponseWriter, r *http.Request) {
		var input CanvasTask
		decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, canvasArtifactResponseLimit+CanvasInputLimit))
		decoder.DisallowUnknownFields()
		if decoder.Decode(&input) != nil || decoder.Decode(new(any)) != io.EOF || len(input.Prompt) > CanvasInputLimit || strings.TrimSpace(input.Prompt) == "" || len(input.Parameters) > 0 || len(input.Images) > 4 {
			fail(w, 400, "invalid_ai_request")
			return
		}
		models, err := workerCanvasModels()
		if err != nil {
			fail(w, 503, "codex_unavailable")
			return
		}
		var selected *CanvasModel
		for _, model := range models {
			if model.ID == input.ModelID {
				copy := model
				selected = &copy
			}
		}
		if selected == nil {
			fail(w, 422, "model_unavailable")
			return
		}
		select {
		case slot <- struct{}{}:
			defer func() { <-slot }()
		default:
			w.Header().Set("Retry-After", "5")
			fail(w, 429, "codex_busy")
			return
		}
		executeCanvasTask(w, r, input, *selected)
	})
	server := &http.Server{Handler: mux, ReadHeaderTimeout: 5 * time.Second, ReadTimeout: 10 * time.Second, WriteTimeout: 195 * time.Second}
	return server.Serve(listener)
}

func execute(w http.ResponseWriter, r *http.Request, prompt string) { executeModel(w, r, prompt, "") }
func executeModel(w http.ResponseWriter, r *http.Request, prompt, model string) {
	ctx, cancel := context.WithTimeout(r.Context(), timeout)
	defer cancel()
	cmd := exec.CommandContext(ctx, executable, "--ask-for-approval", "never",
		"--disable", "shell_tool", "--disable", "apps", "--disable", "plugins",
		"--disable", "multi_agent", "--disable", "hooks", "--disable", "memories",
		"-c", `web_search="disabled"`, "exec", "--json", "--sandbox", "read-only", "--skip-git-repo-check", "--ephemeral", "-")
	if model != "" {
		cmd.Args = append(cmd.Args[:len(cmd.Args)-1], "--model", model, "-")
	}
	cmd.Dir = "/work"
	cmd.Env = []string{"PATH=/usr/local/bin:/usr/bin:/bin", "HOME=/home/node", "CODEX_HOME=/home/node/.codex", "SSL_CERT_FILE=/etc/ssl/certs/ca-certificates.crt"}
	cmd.Stdin = strings.NewReader(prompt)
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	cmd.Cancel = func() error {
		err := syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
		if errors.Is(err, syscall.ESRCH) {
			return os.ErrProcessDone
		}
		return err
	}
	cmd.WaitDelay = 2 * time.Second
	var output outputBuffer
	cmd.Stdout, cmd.Stderr = &output, io.Discard
	err := cmd.Run()
	switch {
	case output.exceeded:
		fail(w, 502, "codex_output_too_large")
	case errors.Is(ctx.Err(), context.DeadlineExceeded):
		fail(w, 504, "codex_timeout")
	case err != nil:
		fail(w, 502, "codex_failed")
	default:
		var result string
		var completed bool
		for _, line := range bytes.Split(bytes.TrimSpace(output.buffer.Bytes()), []byte("\n")) {
			var event struct {
				Type string `json:"type"`
				Item struct {
					Type string `json:"type"`
					Text string `json:"text"`
				} `json:"item"`
			}
			if json.Unmarshal(line, &event) != nil {
				fail(w, 502, "invalid_codex_response")
				return
			}
			if event.Type == "item.completed" && event.Item.Type == "agent_message" {
				result = event.Item.Text
			}
			if event.Type == "turn.failed" || event.Type == "error" {
				fail(w, 502, "codex_failed")
				return
			}
			if event.Type == "turn.completed" {
				completed = true
			}
		}
		if !completed || result == "" {
			fail(w, 502, "codex_failed")
			return
		}
		response, _ := json.Marshal(map[string]string{"result": result})
		if len(response) >= maxOutput {
			fail(w, 502, "codex_output_too_large")
			return
		}
		writeJSON(w, 200, json.RawMessage(response))
	}
}