package main

import (
	"bufio"
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync/atomic"
	"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

var (
	requestLogRoot    = "/var/log/daiapi"
	requestLogCounter uint64
)

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, requestLoggingMiddleware(mux)); err != nil {
		log.Fatal(err)
	}
}

func requestLoggingMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		receivedAt := time.Now()
		body, readErr := io.ReadAll(r.Body)
		if closeErr := r.Body.Close(); closeErr != nil {
			log.Printf("failed to close request body: %v", closeErr)
		}
		if readErr != nil {
			log.Printf("failed to read request body for logging: %v", readErr)
		}
		r.Body = io.NopCloser(bytes.NewReader(body))

		logFile, logPath, logErr := createRequestLogFile(receivedAt, r.URL.Path, body)
		if logErr != nil {
			log.Printf("failed to create request log file: %v", logErr)
		} else {
			defer func() {
				if err := logFile.Close(); err != nil {
					log.Printf("failed to close request log file %s: %v", logPath, err)
				}
			}()
			writeRequestLogStart(logFile, receivedAt, r, body, readErr)
		}

		lrw := &loggingResponseWriter{
			ResponseWriter: w,
			statusCode:     http.StatusOK,
		}

		next.ServeHTTP(lrw, r)

		if logFile != nil {
			writeRequestLogEnd(logFile, time.Now(), lrw.statusCode, lrw.body.Bytes(), nil)
		}
	})
}

func createRequestLogFile(receivedAt time.Time, requestPath string, body []byte) (*os.File, string, error) {
	dir := filepath.Join(
		requestLogRoot,
		receivedAt.Format("2006"),
		receivedAt.Format("01"),
		receivedAt.Format("02"),
	)
	if err := os.MkdirAll(dir, 0750); err != nil {
		return nil, "", err
	}

	for attempts := 0; attempts < 5; attempts++ {
		hash := requestLogHash(receivedAt, requestPath, body)
		name := fmt.Sprintf("%s.%s.log", receivedAt.Format("15-04-05"), hash)
		path := filepath.Join(dir, name)
		file, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0640)
		if err == nil {
			return file, path, nil
		}
		if !errors.Is(err, os.ErrExist) {
			return nil, "", err
		}
	}

	return nil, "", fmt.Errorf("failed to create unique request log file after retries")
}

func requestLogHash(receivedAt time.Time, requestPath string, body []byte) string {
	counter := atomic.AddUint64(&requestLogCounter, 1)
	hash := sha256.New()
	fmt.Fprintf(hash, "%s\n%s\n%d\n", receivedAt.Format(time.RFC3339Nano), requestPath, counter)
	hash.Write(body)
	return hex.EncodeToString(hash.Sum(nil))[:16]
}

func writeRequestLogStart(file *os.File, receivedAt time.Time, r *http.Request, body []byte, readErr error) {
	entry := map[string]interface{}{
		"timestamp": receivedAt.Format(time.RFC3339Nano),
		"method":    r.Method,
		"path":      r.URL.Path,
		"query":     r.URL.Query(),
		"headers":   redactedHeaders(r.Header),
		"body":      string(body),
	}
	if readErr != nil {
		entry["bodyReadError"] = readErr.Error()
	}

	writeLogSection(file, "request", entry)
}

func writeRequestLogEnd(file *os.File, finishedAt time.Time, statusCode int, body []byte, processingErr error) {
	entry := map[string]interface{}{
		"timestamp":    finishedAt.Format(time.RFC3339Nano),
		"statusCode":   statusCode,
		"responseBody": string(body),
	}
	if processingErr != nil {
		entry["error"] = processingErr.Error()
	} else if statusCode >= http.StatusBadRequest {
		entry["error"] = http.StatusText(statusCode)
	}

	writeLogSection(file, "response", entry)
}

func writeLogSection(file *os.File, section string, entry map[string]interface{}) {
	payload, err := json.MarshalIndent(entry, "", "  ")
	if err != nil {
		log.Printf("failed to marshal %s log section: %v", section, err)
		return
	}

	if _, err := fmt.Fprintf(file, "=== %s ===\n%s\n\n", section, payload); err != nil {
		log.Printf("failed to write %s log section: %v", section, err)
	}
}

func redactedHeaders(headers http.Header) http.Header {
	redacted := make(http.Header, len(headers))
	for key, values := range headers {
		if isSensitiveHeader(key) {
			redacted[key] = []string{"[redacted]"}
			continue
		}
		redacted[key] = append([]string(nil), values...)
	}
	return redacted
}

func isSensitiveHeader(key string) bool {
	switch strings.ToLower(key) {
	case "authorization", "cookie", "set-cookie", "x-api-key", "x-auth-token":
		return true
	default:
		return false
	}
}

type loggingResponseWriter struct {
	http.ResponseWriter
	statusCode  int
	wroteHeader bool
	body        bytes.Buffer
}

func (w *loggingResponseWriter) WriteHeader(statusCode int) {
	if w.wroteHeader {
		return
	}
	w.statusCode = statusCode
	w.wroteHeader = true
	w.ResponseWriter.WriteHeader(statusCode)
}

func (w *loggingResponseWriter) Write(data []byte) (int, error) {
	if !w.wroteHeader {
		w.WriteHeader(http.StatusOK)
	}

	n, err := w.ResponseWriter.Write(data)
	if n > 0 {
		w.body.Write(data[:n])
	}
	return n, err
}

func (w *loggingResponseWriter) Flush() {
	if flusher, ok := w.ResponseWriter.(http.Flusher); ok {
		if !w.wroteHeader {
			w.WriteHeader(http.StatusOK)
		}
		flusher.Flush()
	}
}

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)
	}
}