package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/config"
	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/gateway"
	"github.com/tidbcloud/AgentTeam/internal/observability/logging"
	"github.com/tidbcloud/AgentTeam/internal/store/workspace"
)

const workspaceIdentityResolverPath = "/internal/meshxruntime/workspace-route-identity"

func main() {
	if exitCode := handleHelp(os.Args[1:], os.Stdout); exitCode >= 0 {
		os.Exit(exitCode)
	}

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	if err := run(ctx, os.Stderr); err != nil {
		os.Exit(1)
	}
}

func handleHelp(args []string, output io.Writer) int {
	flags := flag.NewFlagSet("agentteam-gateway-server", flag.ContinueOnError)
	flags.SetOutput(output)
	flags.Usage = func() {
		_, _ = fmt.Fprintln(output, "Usage: agentteam-gateway-server [--help]")
		_, _ = fmt.Fprintln(output)
		_, _ = fmt.Fprintln(output, "Runs the SaaS agentteam-gateway-server edge router.")
		_, _ = fmt.Fprintln(output)
		_, _ = fmt.Fprintln(output, "Required environment:")
		_, _ = fmt.Fprintln(output, "  AGENTTEAM_DATABASE_DSN")
		_, _ = fmt.Fprintln(output, "  AGENTTEAM_GATEWAY_CENTRAL_UPSTREAM_URL")
		_, _ = fmt.Fprintln(output, "  AGENTTEAM_GATEWAY_WORKSPACE_UPSTREAMS")
		_, _ = fmt.Fprintln(output, "  AGENTTEAM_INTERNAL_ADMIN_TOKEN")
		_, _ = fmt.Fprintln(output)
		_, _ = fmt.Fprintln(output, "Optional environment:")
		_, _ = fmt.Fprintln(output, "  AGENTTEAM_GATEWAY_WORKSPACE_IDENTITY_RESOLVER_URL")
		_, _ = fmt.Fprintln(output, "    Defaults to the configured default workspace group's upstream plus")
		_, _ = fmt.Fprintln(output, "    /internal/meshxruntime/workspace-route-identity.")
		_, _ = fmt.Fprintln(output, "  AGENTTEAM_GATEWAY_EXTENSION_UPSTREAM_URL")
		_, _ = fmt.Fprintln(output, "    ProfileExtension reverse-proxy target for split deployments. It")
		_, _ = fmt.Fprintln(output, "    must be origin-only (http(s)://host[:port] with no path, query,")
		_, _ = fmt.Fprintln(output, "    userinfo, or fragment); the request path and query are forwarded")
		_, _ = fmt.Fprintln(output, "    unchanged. When unset, /runtime/ext routes fail closed at the")
		_, _ = fmt.Fprintln(output, "    gateway rather than being routed to a central or workspace upstream.")
	}
	flags.Bool("help", false, "show help")
	flags.Bool("h", false, "show help")
	if err := flags.Parse(args); err != nil {
		return 2
	}
	if flags.Lookup("help").Value.String() == "true" || flags.Lookup("h").Value.String() == "true" {
		flags.Usage()
		return 0
	}
	return -1
}

func run(ctx context.Context, bootstrap io.Writer) (returnErr error) {
	cfg, err := config.LoadFromEnv()
	if err != nil {
		logging.WriteBootstrapError(bootstrap, "load config", err)
		return fmt.Errorf("load config: %w", err)
	}
	pipeline, err := logging.Open(cfg.Log, os.Stdout)
	if err != nil {
		logging.WriteBootstrapError(bootstrap, "initialize logging", err)
		return fmt.Errorf("initialize logging: %w", err)
	}
	logger := pipeline.Logger()
	previousLogger := slog.Default()
	slog.SetDefault(logger)
	defer slog.SetDefault(previousLogger)
	defer func() {
		if returnErr != nil {
			logger.Error("run agentteam gateway server", slog.Any("error", returnErr))
		}
		if err := pipeline.Close(); err != nil {
			logging.WriteBootstrapError(bootstrap, "close logging", err)
			returnErr = errors.Join(returnErr, err)
		}
	}()
	defer logging.RecoverPanic(logger)

	admin := coredb.NewAdmin()
	adminDB, err := admin.Open(ctx, cfg.Database.URL, cfg.Database.ConnectTimeout)
	if err != nil {
		return fmt.Errorf("open gateway database: %w", err)
	}
	defer adminDB.Close()
	if err := admin.RequireSchemaCurrent(ctx, adminDB); err != nil {
		return fmt.Errorf("require gateway database schema current: %w", err)
	}

	db, ok := adminDB.(coredb.DB)
	if !ok {
		return fmt.Errorf("gateway database does not implement coredb.DB: %T", adminDB)
	}
	if raw := coredb.RawSQLDB(db); raw != nil {
		raw.SetMaxOpenConns(cfg.Database.Pool.MaxOpenConns)
		raw.SetMaxIdleConns(cfg.Database.Pool.MaxIdleConns)
		raw.SetConnMaxLifetime(cfg.Database.Pool.ConnMaxLifetime)
		raw.SetConnMaxIdleTime(cfg.Database.Pool.ConnMaxIdleTime)
	}

	workspaceUpstreams, err := parseWorkspaceUpstreams(os.Getenv("AGENTTEAM_GATEWAY_WORKSPACE_UPSTREAMS"))
	if err != nil {
		return err
	}
	identityResolverURL, err := resolveWorkspaceIdentityResolverURL(
		os.Getenv("AGENTTEAM_GATEWAY_WORKSPACE_IDENTITY_RESOLVER_URL"),
		cfg.Server.DefaultWorkspaceGroupIDOrDefault(),
		workspaceUpstreams,
	)
	if err != nil {
		return err
	}
	identityResolver, err := gateway.NewHTTPWorkspaceIdentityResolver(
		identityResolverURL,
		cfg.Server.InternalAdminToken,
		nil,
	)
	if err != nil {
		return fmt.Errorf("create gateway workspace identity resolver: %w", err)
	}
	handler, err := gateway.NewHandler(workspace.New(db, workspace.Options{}), gateway.Config{
		CentralUpstream:    strings.TrimSpace(os.Getenv("AGENTTEAM_GATEWAY_CENTRAL_UPSTREAM_URL")),
		WorkspaceUpstreams: workspaceUpstreams,
		IdentityResolver:   identityResolver,
		TrustedJWTHeader:   strings.TrimSpace(os.Getenv("MESHX_AUTH_JWT_HEADER")),
		ExtensionUpstream:  strings.TrimSpace(os.Getenv("AGENTTEAM_GATEWAY_EXTENSION_UPSTREAM_URL")),
	})
	if err != nil {
		return fmt.Errorf("create gateway handler: %w", err)
	}

	server := &http.Server{
		Addr:              cfg.HTTP.Address,
		Handler:           handler,
		ErrorLog:          logging.HTTPErrorLog(logger),
		ReadHeaderTimeout: 5 * time.Second,
	}

	errCh := make(chan error, 1)
	logging.Go(logger, func() {
		logger.Info("agentteam gateway server listening", slog.String("address", cfg.HTTP.Address))
		errCh <- server.ListenAndServe()
	})

	select {
	case err := <-errCh:
		if errors.Is(err, http.ErrServerClosed) {
			return nil
		}
		return err
	case <-ctx.Done():
		shutdownCtx, cancel := context.WithTimeout(context.Background(), cfg.HTTP.ShutdownTimeout)
		defer cancel()
		if err := server.Shutdown(shutdownCtx); err != nil {
			return fmt.Errorf("shutdown gateway http server: %w", err)
		}
		return nil
	}
}

func resolveWorkspaceIdentityResolverURL(explicitURL, defaultWorkspaceGroup string, workspaceUpstreams map[string]string) (string, error) {
	if configured := strings.TrimSpace(explicitURL); configured != "" {
		return configured, nil
	}
	groupID := strings.TrimSpace(defaultWorkspaceGroup)
	upstream := strings.TrimSpace(workspaceUpstreams[groupID])
	if upstream == "" {
		return "", fmt.Errorf("derive gateway workspace identity resolver URL: default workspace group %q has no configured workspace upstream", groupID)
	}
	resolverURL, err := url.JoinPath(upstream, workspaceIdentityResolverPath)
	if err != nil {
		return "", fmt.Errorf("derive gateway workspace identity resolver URL from group %q: %w", groupID, err)
	}
	return resolverURL, nil
}

func parseWorkspaceUpstreams(raw string) (map[string]string, error) {
	result := make(map[string]string)
	for _, item := range strings.Split(raw, ",") {
		item = strings.TrimSpace(item)
		if item == "" {
			continue
		}
		groupID, upstream, ok := strings.Cut(item, "=")
		if !ok {
			return nil, fmt.Errorf("parse AGENTTEAM_GATEWAY_WORKSPACE_UPSTREAMS item %q: expected group=url", item)
		}
		groupID = strings.TrimSpace(groupID)
		upstream = strings.TrimSpace(upstream)
		if groupID == "" || upstream == "" {
			return nil, fmt.Errorf("parse AGENTTEAM_GATEWAY_WORKSPACE_UPSTREAMS item %q: group and url are required", item)
		}
		result[groupID] = upstream
	}
	return result, nil
}