package server

import (
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"

	"github.com/indifferentketchup/llama-sidecar/internal/pool"
)

var shellUnsafe = strings.NewReplacer(
	"`", "", "$", "", "|", "", ";", "", "&", "", "\n", "",
)

func parseFlags(raw string) ([]string, error) {
	cleaned := shellUnsafe.Replace(raw)
	if cleaned != raw {
		return nil, fmt.Errorf("flags contain unsafe characters")
	}
	return splitArgs(strings.TrimSpace(raw)), nil
}

func splitArgs(s string) []string {
	if s == "" {
		return nil
	}
	return strings.Fields(s)
}

func proxyHandler(p *pool.Pool) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		flagsRaw := r.Header.Get("X-Agent-Flags")
		var flags []string
		if flagsRaw != "" {
			var err error
			flags, err = parseFlags(flagsRaw)
			if err != nil {
				writeJSON(w, http.StatusBadRequest, map[string]string{
					"error": err.Error(),
				})
				return
			}
		}

		modelID := r.Header.Get("X-Model-Id")
		if modelID == "" {
			body, err := io.ReadAll(io.LimitReader(r.Body, 1<<20))
			if err != nil {
				writeJSON(w, http.StatusBadRequest, map[string]string{"error": "failed to read body"})
				return
			}
			var req struct {
				Model string `json:"model"`
			}
			if err := json.Unmarshal(body, &req); err == nil && req.Model != "" {
				modelID = req.Model
			}
			r.Body = io.NopCloser(strings.NewReader(string(body)))
			r.ContentLength = int64(len(body))
		}
		if modelID == "" {
			writeJSON(w, http.StatusBadRequest, map[string]string{"error": "model not specified (X-Model-Id header or body.model)"})
			return
		}

		sidecar, err := p.Acquire(r.Context(), modelID, flags)
		if err != nil {
			errMsg := err.Error()
			status := http.StatusInternalServerError
			if strings.Contains(errMsg, "validation:") {
				status = http.StatusBadRequest
			} else if strings.Contains(errMsg, "unknown model:") {
				status = http.StatusNotFound
			} else if strings.Contains(errMsg, "port allocation:") {
				status = http.StatusServiceUnavailable
			}
			writeJSON(w, status, map[string]string{"error": errMsg})
			return
		}

		target := &url.URL{
			Scheme: "http",
			Host:   fmt.Sprintf("127.0.0.1:%d", sidecar.Port),
		}
		proxy := httputil.NewSingleHostReverseProxy(target)
		proxy.ErrorHandler = func(rw http.ResponseWriter, req *http.Request, err error) {
			slog.Error("upstream error", "hash", sidecar.Hash, "port", sidecar.Port, "err", err)
			writeJSON(rw, http.StatusBadGateway, map[string]any{
				"error":        "upstream unavailable",
				"error_detail": err.Error(),
				"sidecar_hash": sidecar.Hash,
				"sidecar_port": sidecar.Port,
				"last_stderr":  sidecar.LastStderr(),
			})
		}

		sidecar.TouchLastUsed()
		proxy.ServeHTTP(w, r)
	}
}

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	json.NewEncoder(w).Encode(v)
}