package utils

import (
	stdcontext "context"
	"encoding/json"
	"net/http"

	"github.com/labstack/echo/v4"
	"gopkg.in/yaml.v3"

	"github.com/pingcap/test-infra/services/tcms/utils/log"
)

const (
	// For JWT middleware.
	CONTEXT_KEY_JWT = "jwt"
	// Context key to store user information from the token into context.
	CONTEXT_KEY_USER    = "user"
	CONTEXT_KEY_DRY_RUN = "dryrun"
)

type Context interface {
	echo.Context

	// current user (Only exist when the request is passed by Auth middleware.)
	User() *SessionUser
	SetUser(user *SessionUser)

	Log() *log.Logger
	ReqID() string
	SetLogReqID(reqID string)

	DryRun() bool
	SetDryRun(dryRun bool)

	JSON(code int, i interface{}) error

	// YAML writes the YAML encoding of v to the response.
	// If multiple items are passed in, the
	// second and subsequent document will be preceded
	// with a "---" document separator, but the first will not.
	YAML(code int, xs ...interface{}) error
}

func NewContext(c echo.Context) echo.Context {
	return &context{
		Context: c,
		log:     log.Log,
	}
}

type context struct {
	echo.Context
	log *log.Logger
}

func (c *context) User() *SessionUser {
	u, ok := c.Get(CONTEXT_KEY_USER).(*SessionUser)
	if ok {
		return u
	}
	return nil
}

func (c *context) SetUser(user *SessionUser) {
	c.Set(CONTEXT_KEY_USER, user)
	c.log = c.log.With("userID", user.ID)
}

func (c *context) ReqID() string {
	return c.Request().Header.Get(echo.HeaderXRequestID)
}

func (c *context) Log() *log.Logger {
	return c.log
}

func (c *context) SetLogReqID(reqID string) {
	c.log = c.log.With("reqID", reqID)
}

func (c *context) SetDryRun(dryRun bool) {
	c.Set(CONTEXT_KEY_DRY_RUN, dryRun)
	c.log = c.log.With("dryrun", dryRun)
}

func (c *context) DryRun() bool {
	if dryRun, ok := c.Get(CONTEXT_KEY_DRY_RUN).(bool); ok {
		return dryRun
	}
	return false
}

func (c *context) JSON(code int, i interface{}) error {
	resp := c.Response()
	resp.Status = code
	resp.Header().Set(echo.HeaderContentType, echo.MIMEApplicationJSONCharsetUTF8)

	e := json.NewEncoder(resp)
	e.SetEscapeHTML(false)
	if _, pretty := c.QueryParams()["pretty"]; pretty {
		e.SetIndent("", "  ")
	}
	return e.Encode(i)
}

func (c *context) YAML(code int, xs ...interface{}) error {
	resp := c.Response()
	resp.Status = code
	resp.Header().Set(echo.HeaderContentType, "application/x-yaml")
	e := yaml.NewEncoder(resp)
	for _, x := range xs {
		if err := e.Encode(x); err != nil {
			return err
		}
	}
	return nil
}

var (
	fakeEcho = echo.New()
)

func FakeCtx() Context {
	fakeRequest, _ := http.NewRequestWithContext(stdcontext.Background(), "GET", "/", nil)
	ctx := NewContext(fakeEcho.NewContext(fakeRequest, nil)).(Context)
	ctx.SetLogReqID(NewReqID())
	ctx.SetUser(&SessionUser{ID: 0})
	ctx.Set(CONTEXT_KEY_JWT, &SessionUser{ID: 0})
	return ctx
}

type ContextKey string

const (
	UtilContextKey ContextKey = "util.Context"
)

func EmbedIntoContext(ctx Context) stdcontext.Context {
	return stdcontext.WithValue(ctx.Request().Context(), UtilContextKey, ctx)
}

func GetUtilContext(ctx stdcontext.Context) Context {
	return ctx.Value(UtilContextKey).(Context)
}