package main

import (
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
	"time"

	"secure-ollama-api/utils"

	"git.secnex.io/secnex/masterlog"
)

type Proxy struct {
	ReverseProxy *httputil.ReverseProxy
	TargetURL    *url.URL
	APIKey       string
}

type responseWriter struct {
	http.ResponseWriter
	statusCode int
}

func NewProxy() (*Proxy, error) {
	remote, err := url.Parse(utils.GetEnv("OLLAMA_SERVER_URL", "http://ollama:11434"))
	if err != nil {
		masterlog.Error(err.Error())
		return nil, err
	}

	apiKey := utils.GetEnv("OLLAMA_API_KEY", "")
	if apiKey == "" {
		masterlog.Error("OLLAMA_API_KEY environment variable is not set")
	}

	return &Proxy{
		ReverseProxy: httputil.NewSingleHostReverseProxy(remote),
		TargetURL:    remote,
		APIKey:       apiKey,
	}, nil
}

func (rw *responseWriter) WriteHeader(code int) {
	rw.statusCode = code
	rw.ResponseWriter.WriteHeader(code)
}

func (p *Proxy) RequestHandler() func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		masterlog.Info("Request received", map[string]interface{}{
			"url": r.URL.String(),
			"ip":  r.RemoteAddr,
		})
		startTime := time.Now()

		// Check Authorization header
		authHeader := r.Header.Get("Authorization")
		expectedAuth := "Bearer " + p.APIKey

		if p.APIKey == "" {
			masterlog.Error("OLLAMA_API_KEY is not configured", map[string]interface{}{
				"url": r.URL.String(),
				"ip":  r.RemoteAddr,
			})
			http.Error(w, "Internal server error", http.StatusInternalServerError)
			return
		}

		if authHeader == "" || !strings.HasPrefix(authHeader, "Bearer ") {
			masterlog.Info("Unauthorized request - missing or invalid Authorization header", map[string]interface{}{
				"url": r.URL.String(),
				"ip":  r.RemoteAddr,
			})
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}

		if authHeader != expectedAuth {
			masterlog.Info("Unauthorized request - invalid API key", map[string]interface{}{
				"url": r.URL.String(),
				"ip":  r.RemoteAddr,
			})
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}

		// Wrap response writer to capture status code
		rw := &responseWriter{
			ResponseWriter: w,
			statusCode:     http.StatusOK, // default status code
		}

		p.ReverseProxy.ServeHTTP(rw, r)

		duration := time.Since(startTime)
		masterlog.Info("Request completed", map[string]interface{}{
			"url":        r.URL.String(),
			"duration":   duration.Milliseconds(),
			"statusCode": rw.statusCode,
			"ip":         r.RemoteAddr,
		})
	}
}

func (p *Proxy) Start() {
	http.HandleFunc("/", p.RequestHandler())
	http.ListenAndServe(":8080", nil)
}