// Package router contains the Phase 2 workspace group routing foundation.
package router

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"

	"github.com/gin-gonic/gin"

	"github.com/tidbcloud/AgentTeam/internal/api"
	"github.com/tidbcloud/AgentTeam/internal/store/workspace"
)

const HeaderResolvedWorkspaceGroupID = "X-AgentTeam-Resolved-Workspace-Group-Id"

// Resolver resolves a workspace to the group that currently owns new traffic.
type Resolver interface {
	ResolveWorkspaceRouteTarget(ctx context.Context, workspaceID string) (workspace.WorkspaceRouteTarget, error)
}

// Handler forwards workspace-scoped traffic to the current owner group.
type Handler struct {
	resolver  Resolver
	upstreams map[string]*url.URL
	proxies   map[string]*httputil.ReverseProxy
}

// newWorkspaceProxy builds one reverse proxy for a workspace group upstream. It
// uses a dedicated short-lived connection pool per proxy: a shared default
// transport would let one upstream's shutdown (CloseIdleConnections) race with
// another proxy's in-flight request and surface as a spurious 502.
func newWorkspaceProxy(upstream *url.URL) *httputil.ReverseProxy {
	proxy := httputil.NewSingleHostReverseProxy(upstream)
	proxy.Transport = &http.Transport{
		Proxy:             http.ProxyFromEnvironment,
		DisableKeepAlives: true,
	}
	originalDirector := proxy.Director
	proxy.Director = func(out *http.Request) {
		originalDirector(out)
		out.Host = upstream.Host
	}
	return proxy
}

// NewHandler creates a workspace router over stable group service upstreams.
func NewHandler(resolver Resolver, upstreams map[string]string) (*Handler, error) {
	if resolver == nil {
		return nil, fmt.Errorf("workspace router resolver is required")
	}
	parsed := make(map[string]*url.URL, len(upstreams))
	proxies := make(map[string]*httputil.ReverseProxy, len(upstreams))
	for groupID, raw := range upstreams {
		trimmedGroupID := strings.TrimSpace(groupID)
		if trimmedGroupID == "" {
			return nil, fmt.Errorf("workspace router upstream group id is required")
		}
		trimmedRaw := strings.TrimSpace(raw)
		if trimmedRaw == "" {
			return nil, fmt.Errorf("workspace router upstream URL is required for group %s", trimmedGroupID)
		}
		upstream, err := url.Parse(trimmedRaw)
		if err != nil {
			return nil, fmt.Errorf("parse workspace router upstream for group %s: %w", trimmedGroupID, err)
		}
		if upstream.Scheme == "" || upstream.Host == "" {
			return nil, fmt.Errorf("workspace router upstream for group %s must include scheme and host", trimmedGroupID)
		}
		parsed[trimmedGroupID] = upstream
		proxies[trimmedGroupID] = newWorkspaceProxy(upstream)
	}
	return &Handler{resolver: resolver, upstreams: parsed, proxies: proxies}, nil
}

// ServeHTTP resolves X-AgentTeam-Route-Workspace-Id and reverse-proxies the request.
func (handler *Handler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	if handler == nil || handler.resolver == nil {
		http.Error(writer, "workspace router is not configured", http.StatusServiceUnavailable)
		return
	}
	workspaceID := routeWorkspaceID(request)
	if workspaceID == "" {
		http.Error(writer, "X-AgentTeam-Route-Workspace-Id is required", http.StatusBadRequest)
		return
	}
	target, err := handler.resolver.ResolveWorkspaceRouteTarget(request.Context(), workspaceID)
	if err != nil {
		status := api.StatusForError(err)
		http.Error(writer, err.Error(), status)
		return
	}
	if resolvedWorkspaceID := strings.TrimSpace(target.WorkspaceID); resolvedWorkspaceID == "" || resolvedWorkspaceID != workspaceID {
		http.Error(writer, "workspace placement identity does not match route workspace", http.StatusConflict)
		return
	}
	upstream, ok := handler.upstreamForTarget(target)
	if !ok {
		http.Error(writer, fmt.Sprintf("workspace group %s has no configured upstream", target.WorkspaceGroupID), http.StatusServiceUnavailable)
		return
	}
	proxy, ok := handler.proxyForTarget(target)
	if !ok {
		http.Error(writer, fmt.Sprintf("workspace group %s has no configured proxy", target.WorkspaceGroupID), http.StatusServiceUnavailable)
		return
	}
	originalDirector := proxy.Director
	proxy.Director = func(out *http.Request) {
		originalDirector(out)
		out.Host = upstream.Host
		out.Header.Set(api.HeaderRouteWorkspaceID, target.WorkspaceID)
		out.Header.Set(api.HeaderWorkspaceID, target.WorkspaceID)
		out.Header.Set(HeaderResolvedWorkspaceGroupID, target.WorkspaceGroupID)
	}
	proxy.ServeHTTP(writer, request)
}

// GinHandler adapts the router to Gin when the deployment embeds it in AgentTeam.
func (handler *Handler) GinHandler() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		handler.ServeHTTP(ctx.Writer, ctx.Request)
	}
}

func (handler *Handler) upstreamForTarget(target workspace.WorkspaceRouteTarget) (*url.URL, bool) {
	for _, key := range []string{target.ServiceRef, target.WorkspaceGroupID} {
		trimmed := strings.TrimSpace(key)
		if trimmed == "" {
			continue
		}
		upstream, ok := handler.upstreams[trimmed]
		if ok {
			return upstream, true
		}
	}
	return nil, false
}

func (handler *Handler) proxyForTarget(target workspace.WorkspaceRouteTarget) (*httputil.ReverseProxy, bool) {
	for _, key := range []string{target.ServiceRef, target.WorkspaceGroupID} {
		trimmed := strings.TrimSpace(key)
		if trimmed == "" {
			continue
		}
		proxy, ok := handler.proxies[trimmed]
		if ok {
			return proxy, true
		}
	}
	return nil, false
}

func routeWorkspaceID(request *http.Request) string {
	if request == nil {
		return ""
	}
	if id := strings.TrimSpace(request.Header.Get(api.HeaderRouteWorkspaceID)); id != "" {
		return id
	}
	// Fallback for older clients that send X-Workspace-Id but not the
	// edge routing hint. This keeps workspace-scoped API calls routable
	// through the gateway even when the client was built before the
	// X-AgentTeam-Route-Workspace-Id header was introduced.
	return strings.TrimSpace(request.Header.Get(api.HeaderWorkspaceID))
}