// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file is distributed
// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
// express or implied. See the License for the specific language governing
// permissions and limitations under the License.

// Package handlers deals with the agent introspection api.
package handlers

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/pprof"
	"strconv"
	"time"

	"github.com/aws/amazon-ecs-agent/agent/config"
	"github.com/aws/amazon-ecs-agent/agent/engine"
	handlersutils "github.com/aws/amazon-ecs-agent/agent/handlers/utils"
	v1 "github.com/aws/amazon-ecs-agent/agent/handlers/v1"
	logginghandler "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/logging"
	"github.com/aws/amazon-ecs-agent/ecs-agent/utils/retry"
	"github.com/cihub/seelog"
)

type rootResponse struct {
	AvailableCommands []string
}

const (
	// With pprof we need to increase the timeout so that there is enough time to do the profiling. Since the profiling
	// time window for CPU is configurable in the request, this timeout effectively means the CPU profiling will be
	// capped to 5 min.
	writeTimeoutForPprof = time.Minute * 5
	pprofBasePath        = "/debug/pprof/"
	pprofCMDLinePath     = pprofBasePath + "cmdline"
	pprofProfilePath     = pprofBasePath + "profile"
	pprofSymbolPath      = pprofBasePath + "symbol"
	pprofTracePath       = pprofBasePath + "trace"
)

var (
	// Injection points for testing
	pprofIndexHandler   = pprof.Index
	pprofCmdlineHandler = pprof.Cmdline
	pprofProfileHandler = pprof.Profile
	pprofSymbolHandler  = pprof.Symbol
	pprofTraceHandler   = pprof.Trace
)

func introspectionServerSetup(containerInstanceArn *string, taskEngine handlersutils.DockerStateResolver, cfg *config.Config) *http.Server {
	paths := []string{v1.AgentMetadataPath, v1.TaskContainerMetadataPath, v1.LicensePath}

	if cfg.EnableRuntimeStats.Enabled() {
		paths = append(paths, pprofBasePath, pprofCMDLinePath, pprofProfilePath, pprofSymbolPath, pprofTracePath)
	}

	availableCommands := &rootResponse{paths}
	// Autogenerated list of the above serverFunctions paths
	availableCommandResponse, err := json.Marshal(&availableCommands)
	if err != nil {
		seelog.Errorf("Error marshaling JSON in introspection server setup: %s", err)
	}

	defaultHandler := func(w http.ResponseWriter, r *http.Request) {
		w.Write(availableCommandResponse)
	}

	serverMux := http.NewServeMux()
	serverMux.HandleFunc("/", defaultHandler)

	v1HandlersSetup(serverMux, containerInstanceArn, taskEngine, cfg)
	pprofHandlerSetup(serverMux, cfg)

	// Log all requests and then pass through to serverMux
	loggingServeMux := http.NewServeMux()
	loggingServeMux.Handle("/", logginghandler.NewLoggingHandler(serverMux))

	wTimeout := writeTimeout
	if cfg.EnableRuntimeStats.Enabled() {
		wTimeout = writeTimeoutForPprof
	}
	server := &http.Server{
		Addr:         ":" + strconv.Itoa(config.AgentIntrospectionPort),
		Handler:      loggingServeMux,
		ReadTimeout:  readTimeout,
		WriteTimeout: wTimeout,
	}

	return server
}

// v1HandlersSetup adds all handlers except CredentialsHandler in v1 package to the server mux.
func v1HandlersSetup(serverMux *http.ServeMux,
	containerInstanceArn *string,
	taskEngine handlersutils.DockerStateResolver,
	cfg *config.Config) {
	serverMux.HandleFunc(v1.AgentMetadataPath, v1.AgentMetadataHandler(containerInstanceArn, cfg))
	serverMux.HandleFunc(v1.TaskContainerMetadataPath, v1.TaskContainerMetadataHandler(taskEngine))
	serverMux.HandleFunc(v1.LicensePath, v1.LicenseHandler)
}

func pprofHandlerSetup(serverMux *http.ServeMux, cfg *config.Config) {
	if !cfg.EnableRuntimeStats.Enabled() {
		return
	}
	serverMux.HandleFunc(pprofBasePath, pprofIndexHandler)
	serverMux.HandleFunc(pprofCMDLinePath, pprofCmdlineHandler)
	serverMux.HandleFunc(pprofProfilePath, pprofProfileHandler)
	serverMux.HandleFunc(pprofSymbolPath, pprofSymbolHandler)
	serverMux.HandleFunc(pprofTracePath, pprofTraceHandler)
}

// ServeIntrospectionHTTPEndpoint serves information about this agent/containerInstance and tasks
// running on it. "V1" here indicates the hostname version of this server instead
// of the handler versions, i.e. "V1" server can include "V1" and "V2" handlers.
func ServeIntrospectionHTTPEndpoint(ctx context.Context, containerInstanceArn *string, taskEngine engine.TaskEngine, cfg *config.Config) {
	// Is this the right level to type assert, assuming we'd abstract multiple taskengines here?
	// Revisit if we ever add another type..
	dockerTaskEngine := taskEngine.(*engine.DockerTaskEngine)

	server := introspectionServerSetup(containerInstanceArn, dockerTaskEngine, cfg)

	go func() {
		<-ctx.Done()
		if err := server.Shutdown(context.Background()); err != nil {
			// Error from closing listeners, or context timeout:
			seelog.Infof("HTTP server Shutdown: %v", err)
		}
	}()

	for {
		retry.RetryWithBackoff(retry.NewExponentialBackoff(time.Second, time.Minute, 0.2, 2), func() error {
			if err := server.ListenAndServe(); err != http.ErrServerClosed {
				seelog.Errorf("Error running introspection endpoint: %v", err)
				return err
			}
			// server was cleanly closed via context
			return nil
		})
	}
}