package main

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

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

type runResponse struct {
	Success bool        `json:"success"`
	Answer  *string     `json:"answer"`
	Usage   interface{} `json:"usage"`
	Stderr  string      `json:"stderr"`
}

type geminiResponse struct {
	Success  bool        `json:"success"`
	Answer   *string     `json:"answer"`
	Usage    interface{} `json:"usage"`
	Stdout   string      `json:"stdout"`
	Stderr   string      `json:"stderr"`
	ExitCode int         `json:"exitCode"`
	Error    string      `json:"error,omitempty"`
}

type event struct {
	Type  string                 `json:"type"`
	Item  map[string]interface{} `json:"item"`
	Usage interface{}            `json:"usage"`
}

type commandResult struct {
	Stdout   string
	Stderr   string
	ExitCode int
	Err      error
	TimedOut bool
}

const geminiTimeout = 5 * time.Minute

func main() {
	mux := http.NewServeMux()
	mux.HandleFunc("/run", runCodexHandler)
	mux.HandleFunc("/runCodex", runCodexHandler)
	mux.HandleFunc("/runGemini", runGeminiHandler)

	addr := ":8000"
	if port := os.Getenv("PORT"); port != "" {
		addr = ":" + port
	}

	log.Printf("listening on %s", addr)
	if err := http.ListenAndServe(addr, mux); err != nil {
		log.Fatal(err)
	}
}

func runCodexHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	var req runRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "invalid json body", http.StatusBadRequest)
		return
	}

	cmd := exec.Command(
		"codex",
		"exec",
		"--skip-git-repo-check",
		"--full-auto",
		"--add-dir", "/tmp",
		"--json",
		req.Prompt,
	)

	var stdout bytes.Buffer
	var stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	err := cmd.Run()

	var finalText *string
	var usage interface{}

	scanner := bufio.NewScanner(bytes.NewReader(stdout.Bytes()))
	for scanner.Scan() {
		line := bytes.TrimSpace(scanner.Bytes())
		if len(line) == 0 {
			continue
		}

		var evt event
		if unmarshalErr := json.Unmarshal(line, &evt); unmarshalErr != nil {
			continue
		}

		if evt.Type == "item.completed" {
			itemType, _ := evt.Item["type"].(string)
			itemText, _ := evt.Item["text"].(string)
			if (itemType == "agent_message" || itemType == "message") && itemText != "" {
				textCopy := itemText
				finalText = &textCopy
			}
		}

		if evt.Type == "turn.completed" {
			usage = evt.Usage
		}
	}

	resp := runResponse{
		Success: err == nil,
		Answer:  finalText,
		Usage:   usage,
		Stderr:  stderr.String(),
	}

	w.Header().Set("Content-Type", "application/json")
	if encodeErr := json.NewEncoder(w).Encode(resp); encodeErr != nil {
		http.Error(w, encodeErr.Error(), http.StatusInternalServerError)
	}
}

func runGeminiHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")

	if r.Method != http.MethodPost {
		writeGeminiError(w, http.StatusMethodNotAllowed, "method not allowed")
		return
	}

	var req runRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeGeminiError(w, http.StatusBadRequest, "invalid json body")
		return
	}

	prompt := strings.TrimSpace(req.Prompt)
	if prompt == "" {
		writeGeminiError(w, http.StatusBadRequest, "prompt must not be empty")
		return
	}

	result := runCommand(r.Context(), geminiTimeout, "gemini", "-p", req.Prompt)
	answer := result.Stdout
	resp := geminiResponse{
		Success:  result.Err == nil,
		Answer:   &answer,
		Usage:    nil,
		Stdout:   result.Stdout,
		Stderr:   result.Stderr,
		ExitCode: result.ExitCode,
		Error:    commandError(result),
	}

	if result.Err != nil {
		w.WriteHeader(http.StatusInternalServerError)
	}

	if encodeErr := json.NewEncoder(w).Encode(resp); encodeErr != nil {
		http.Error(w, encodeErr.Error(), http.StatusInternalServerError)
	}
}

func runCommand(parent context.Context, timeout time.Duration, name string, args ...string) commandResult {
	ctx, cancel := context.WithTimeout(parent, timeout)
	defer cancel()

	cmd := exec.CommandContext(ctx, name, args...)

	var stdout bytes.Buffer
	var stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	err := cmd.Run()
	result := commandResult{
		Stdout:   stdout.String(),
		Stderr:   stderr.String(),
		ExitCode: 0,
		Err:      err,
		TimedOut: ctx.Err() == context.DeadlineExceeded,
	}

	if err != nil {
		result.ExitCode = -1
		var exitErr *exec.ExitError
		if errors.As(err, &exitErr) {
			result.ExitCode = exitErr.ExitCode()
		}
	}

	return result
}

func commandError(result commandResult) string {
	if result.TimedOut {
		return "process timed out"
	}
	if result.Err == nil {
		return ""
	}
	if errors.Is(result.Err, exec.ErrNotFound) {
		return "gemini CLI not found in PATH"
	}

	var execErr *exec.Error
	if errors.As(result.Err, &execErr) {
		return "gemini CLI not found in PATH"
	}

	return result.Err.Error()
}

func writeGeminiError(w http.ResponseWriter, status int, message string) {
	w.WriteHeader(status)
	resp := geminiResponse{
		Success:  false,
		Answer:   nil,
		Usage:    nil,
		Stdout:   "",
		Stderr:   "",
		ExitCode: -1,
		Error:    message,
	}
	if err := json.NewEncoder(w).Encode(resp); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}