package agent

import (
	"bytes"
	"context"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"
)

func (bridge *AgentTeamMaaSLLMBridge) serveHTTP(writer http.ResponseWriter, request *http.Request) {
	startedAt := time.Now()
	if !bridge.authorizeLocalRequest(request) {
		writer.Header().Set("Content-Type", "application/json")
		writer.WriteHeader(http.StatusUnauthorized)
		_, _ = writer.Write([]byte(`{"error":"invalid agentteam llm bridge credential"}`))
		return
	}

	bridge.debugLog("agentteam maas llm bridge local request received",
		"method", request.Method,
		"path", request.URL.Path,
		"query_present", strings.TrimSpace(request.URL.RawQuery) != "",
		"content_length", request.ContentLength,
	)
	body, err := io.ReadAll(request.Body)
	if err != nil {
		bridge.writeBridgeError(writer, http.StatusBadRequest, "read agentteam llm bridge request body")
		return
	}

	headers := cloneHeader(request.Header)
	clearLocalProxyCredentialHeaders(headers)
	if maasGenerationRequest(request.Method, request.URL.Path) && headers.Get("Idempotency-Key") == "" {
		idempotencyKey, keyErr := randomHex(16)
		if keyErr != nil {
			bridge.writeBridgeError(writer, http.StatusBadGateway, "generate maas idempotency key")
			return
		}
		headers.Set("Idempotency-Key", "agentteam-"+idempotencyKey)
	}

	response, err := bridge.doUpstream(request.Context(), request.Method, request.URL.Path, request.URL.RawQuery, headers, body, false)
	if err != nil {
		bridge.writeBridgeError(writer, http.StatusBadGateway, err.Error())
		return
	}
	if response.StatusCode == http.StatusUnauthorized {
		bridge.debugLog("agentteam maas llm bridge upstream unauthorized; refreshing token", "method", request.Method, "path", request.URL.Path)
		_ = response.Body.Close()
		response, err = bridge.doUpstream(request.Context(), request.Method, request.URL.Path, request.URL.RawQuery, headers, body, true)
		if err != nil {
			bridge.writeBridgeError(writer, http.StatusBadGateway, err.Error())
			return
		}
	}
	defer response.Body.Close()
	bridge.debugLog("agentteam maas llm bridge upstream response ready",
		"method", request.Method,
		"path", request.URL.Path,
		"status", response.StatusCode,
		"duration", time.Since(startedAt).Round(time.Millisecond).String(),
	)

	copyHeaders(writer.Header(), response.Header)
	writer.WriteHeader(response.StatusCode)
	copyDestination := io.Writer(writer)
	if flusher, ok := writer.(http.Flusher); ok {
		copyDestination = llmBridgeFlushingWriter{writer: writer, flusher: flusher}
	}
	written, err := io.Copy(copyDestination, response.Body)
	if err != nil {
		// The response may already contain partial model output, so writing a
		// replacement JSON error would corrupt the stream. Close the downstream
		// connection and let the runtime observe the upstream transport failure.
		bridge.debugLog("agentteam maas llm bridge response copy failed",
			"method", request.Method,
			"path", request.URL.Path,
			"bytes", written,
			"duration", time.Since(startedAt).Round(time.Millisecond).String(),
			"error", err,
		)
		panic(http.ErrAbortHandler)
	}
	bridge.debugLog("agentteam maas llm bridge response completed",
		"method", request.Method,
		"path", request.URL.Path,
		"status", response.StatusCode,
		"bytes", written,
		"duration", time.Since(startedAt).Round(time.Millisecond).String(),
	)
}

type llmBridgeFlushingWriter struct {
	writer  io.Writer
	flusher http.Flusher
}

func (writer llmBridgeFlushingWriter) Write(data []byte) (int, error) {
	written, err := writer.writer.Write(data)
	if written > 0 {
		writer.flusher.Flush()
	}
	return written, err
}

func (bridge *AgentTeamMaaSLLMBridge) doUpstream(ctx context.Context, method string, requestPath string, rawQuery string, headers http.Header, body []byte, forceRefresh bool) (*http.Response, error) {
	startedAt := time.Now()
	token, err := bridge.ensureMaaSProxyToken(ctx, forceRefresh)
	if err != nil {
		return nil, fmt.Errorf("prepare maas runtime capability: %w", err)
	}

	upstreamURL := *bridge.upstreamBaseURL
	upstreamURL.Path = requestPath
	upstreamURL.RawPath = requestPath
	upstreamURL.RawQuery = rawQuery

	req, err := http.NewRequestWithContext(ctx, method, upstreamURL.String(), bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("build upstream request: %w", err)
	}
	req.Header = cloneHeader(headers)
	req.Header.Set("Authorization", "Bearer "+token)
	bridge.debugLog("agentteam maas llm bridge upstream request sending",
		"method", method,
		"path", requestPath,
		"query_present", strings.TrimSpace(rawQuery) != "",
		"force_refresh", forceRefresh,
		"upstream_url", sanitizeOpenCodeLogURL(upstreamURL.String()),
		"body_bytes", len(body),
	)

	resp, err := bridge.httpClient.Do(req)
	if err != nil {
		bridge.debugLog("agentteam maas llm bridge upstream request failed",
			"method", method,
			"path", requestPath,
			"force_refresh", forceRefresh,
			"duration", time.Since(startedAt).Round(time.Millisecond).String(),
			"error", err,
		)
		return nil, fmt.Errorf("send upstream request: %w", err)
	}
	bridge.debugLog("agentteam maas llm bridge upstream first byte",
		"method", method,
		"path", requestPath,
		"force_refresh", forceRefresh,
		"status", resp.StatusCode,
		"duration", time.Since(startedAt).Round(time.Millisecond).String(),
	)
	return resp, nil
}

func (bridge *AgentTeamMaaSLLMBridge) authorizeLocalRequest(request *http.Request) bool {
	if bridge == nil || request == nil {
		return false
	}
	return strings.TrimSpace(extractProxyToken(request)) == strings.TrimSpace(bridge.localAPIKey)
}

func (bridge *AgentTeamMaaSLLMBridge) writeBridgeError(writer http.ResponseWriter, status int, message string) {
	writer.Header().Set("Content-Type", "application/json")
	writer.WriteHeader(status)
	_, _ = writer.Write([]byte(fmt.Sprintf(`{"error":%q}`, strings.TrimSpace(message))))
}

func maasGenerationRequest(method string, requestPath string) bool {
	if method != http.MethodPost {
		return false
	}
	normalizedPath := strings.TrimRight(strings.TrimSpace(requestPath), "/")
	switch {
	case strings.HasSuffix(normalizedPath, "/chat/completions"):
		return true
	case strings.HasSuffix(normalizedPath, "/messages"):
		return true
	case strings.HasSuffix(normalizedPath, "/responses"):
		return true
	case strings.HasSuffix(normalizedPath, "/embeddings"):
		return true
	default:
		return false
	}
}

func cloneHeader(headers http.Header) http.Header {
	cloned := make(http.Header, len(headers))
	for key, values := range headers {
		cloned[key] = append([]string(nil), values...)
	}
	return cloned
}

func clearLocalProxyCredentialHeaders(headers http.Header) {
	for key := range headers {
		if reservedAgentTeamBridgeUpstreamHeader(key) {
			headers.Del(key)
		}
	}
}

func reservedAgentTeamBridgeUpstreamHeader(key string) bool {
	normalizedHeader := http.CanonicalHeaderKey(key)
	switch normalizedHeader {
	case "Authorization", "X-Api-Key", "Api-Key",
		"Agentauthtoken", "Agent-Auth-Token", "X-Agent-Auth-Token":
		return true
	default:
		return strings.HasPrefix(normalizedHeader, "X-Agentteam-Maas-")
	}
}

func copyHeaders(dst http.Header, src http.Header) {
	for key, values := range src {
		dst[key] = append([]string(nil), values...)
	}
}

func extractProxyToken(request *http.Request) string {
	if request == nil {
		return ""
	}
	if authorization := strings.TrimSpace(request.Header.Get("Authorization")); authorization != "" {
		if token, ok := strings.CutPrefix(authorization, "Bearer "); ok {
			return strings.TrimSpace(token)
		}
		if token, ok := strings.CutPrefix(authorization, "bearer "); ok {
			return strings.TrimSpace(token)
		}
	}
	if token := strings.TrimSpace(request.Header.Get("X-Api-Key")); token != "" {
		return token
	}
	return strings.TrimSpace(request.Header.Get("Api-Key"))
}

func randomHex(size int) (string, error) {
	if size <= 0 {
		return "", fmt.Errorf("random hex size must be positive")
	}
	buf := make([]byte, size)
	if _, err := rand.Read(buf); err != nil {
		return "", err
	}
	return hex.EncodeToString(buf), nil
}