// Package proxyserver forwards only the proxy extension's two public routes.
package proxyserver

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/subtle"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/http/httputil"
	"os"
	"path/filepath"
	"time"

	"albedo/cli/internal/config"
	"albedo/cli/internal/webui"
)

func KeyPath() string { return filepath.Join(config.HomeDir(), "proxy-key") }

// LoadKey creates or reads the home's private proxy key; rotate replaces it atomically.
func LoadKey(rotate bool) (string, error) {
	path := KeyPath()
	if !rotate {
		key, err := readKey()
		if err == nil {
			return key, os.Chmod(path, 0o600)
		}
		if !errors.Is(err, os.ErrNotExist) {
			return "", err
		}
	}
	raw := make([]byte, 24)
	if _, err := rand.Read(raw); err != nil {
		return "", err
	}
	key := hex.EncodeToString(raw)
	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
		return "", err
	}
	if !rotate {
		f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
		if errors.Is(err, os.ErrExist) {
			return readKey()
		}
		if err != nil {
			return "", err
		}
		_, err = f.WriteString(key)
		return key, errors.Join(err, f.Close())
	}
	f, err := os.CreateTemp(filepath.Dir(path), ".proxy-key-*")
	if err != nil {
		return "", err
	}
	defer os.Remove(f.Name())
	_, err = f.WriteString(key)
	if err = errors.Join(err, f.Close()); err != nil {
		return "", err
	}
	return key, os.Rename(f.Name(), path)
}

func readKey() (string, error) {
	raw, err := os.ReadFile(KeyPath())
	if err != nil {
		return "", err
	}
	key := string(raw)
	if _, err := hex.DecodeString(key); err != nil || len(key) != 48 {
		return "", fmt.Errorf("invalid proxy key in %s; use albedo proxy --new-key", KeyPath())
	}
	return key, nil
}

type Server struct {
	*webui.Access
	upstream webui.Upstream
	proxy    *httputil.ReverseProxy
	logger   *log.Logger
}

func New(upstream webui.Upstream, output io.Writer) *Server {
	s := &Server{Access: webui.NewAccess(), upstream: upstream}
	if output != nil {
		s.logger = log.New(output, "", 0)
	}
	s.proxy = &httputil.ReverseProxy{
		Rewrite: func(pr *httputil.ProxyRequest) {
			endpoint := s.upstream.Snapshot()
			pr.SetURL(webui.Loopback(endpoint.Port))
			pr.Out.Host = pr.Out.URL.Host
			pr.Out.URL.Path = "/extensions/proxy" + pr.In.URL.Path
			pr.Out.URL.RawPath = ""
			for _, header := range []string{"Origin", "Cookie", "Referer", "Forwarded", "X-Forwarded-For", "X-Forwarded-Host", "X-Forwarded-Proto"} {
				pr.Out.Header.Del(header)
			}
			pr.Out.Header.Set("Authorization", "Bearer "+endpoint.Token)
		},
		FlushInterval: -1,
		ModifyResponse: func(response *http.Response) error {
			response.Header.Del("Set-Cookie")
			response.Header.Del("Authorization")
			if response.StatusCode == http.StatusUnauthorized {
				return errors.New("daemon endpoint changed")
			}
			return nil
		},
		ErrorHandler: func(w http.ResponseWriter, r *http.Request, _ error) {
			if r.Context().Err() != nil {
				return
			}
			s.refresh()
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(http.StatusServiceUnavailable)
			_ = json.NewEncoder(w).Encode(map[string]map[string]string{"error": {"message": "the albedo daemon isn't running on this machine. start it by running albedo there, then retry.", "type": "service_unavailable"}})
		},
	}
	return s
}

func (s *Server) refresh() {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_ = s.upstream.Refresh(ctx)
}

func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	started := time.Now()
	response := &statusWriter{ResponseWriter: w, status: http.StatusOK}
	model := "-"
	if s.logger != nil {
		defer func() {
			s.logger.Printf("%q %d %s", model, response.status, time.Since(started).Round(time.Millisecond))
		}()
	}
	if !s.Accept(r) {
		http.Error(response, "host, origin or network is not allowed", http.StatusForbidden)
		return
	}
	if r.URL.EscapedPath() != r.URL.Path || (r.URL.Path != "/v1/models" && r.URL.Path != "/v1/chat/completions") {
		http.NotFound(response, r)
		return
	}
	method := http.MethodGet
	if r.URL.Path == "/v1/chat/completions" {
		method = http.MethodPost
	}
	if r.Method != method {
		response.Header().Set("Allow", method)
		response.WriteHeader(http.StatusMethodNotAllowed)
		return
	}
	key, err := readKey()
	if err != nil || subtle.ConstantTimeCompare([]byte(r.Header.Get("Authorization")), []byte("Bearer "+key)) != 1 {
		response.Header().Set("WWW-Authenticate", "Bearer")
		http.Error(response, "proxy API key required", http.StatusUnauthorized)
		return
	}
	if r.Method == http.MethodPost && s.logger != nil {
		body, err := io.ReadAll(http.MaxBytesReader(response, r.Body, 32_000_000))
		if err != nil {
			code := http.StatusBadRequest
			if _, oversized := errors.AsType[*http.MaxBytesError](err); oversized {
				code = http.StatusRequestEntityTooLarge
			}
			http.Error(response, "could not read request body", code)
			return
		}
		var request struct {
			Model string `json:"model"`
		}
		if json.Unmarshal(body, &request) == nil {
			model = request.Model[:min(len(request.Model), 256)]
		}
		r.Body = io.NopCloser(bytes.NewReader(body))
	}
	s.proxy.ServeHTTP(response, r)
}

type statusWriter struct {
	http.ResponseWriter
	status int
}

func (w *statusWriter) WriteHeader(status int) {
	w.status = status
	w.ResponseWriter.WriteHeader(status)
}
func (w *statusWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }