package resource

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	commM "github.com/pingcap/test-infra/common/model"

	"github.com/pingcap/test-infra/services/tps/pkg/common"

	rmsReq "github.com/pingcap/test-infra/common/model/rms/request"
	rmsResp "github.com/pingcap/test-infra/common/model/rms/response"
	"k8s.io/klog/v2"
)

type JobContext struct {
	context.Context
	job *common.ArgoJob
}

func NewJobContext(ctx context.Context, job *common.ArgoJob) JobContext {
	return JobContext{Context: ctx, job: job}
}

func (j JobContext) GetJob() *common.ArgoJob {
	return j.job
}

func NewTcmsResourcePoolMgr() ResourcePoolMgr {
	t := ResourcePoolMgr{poolMetaProvider: TcmsPoolMetaProvider{tcmsClient: Tcms}}
	t.init()
	return t
}

type PoolMeta struct {
	id             commM.ID
	name           string
	alias          string
	rmsUrl         string
	borrowedPolicy commM.ResourcePoolBorrowedPolicies
}

type PoolMetaProvider interface {
	Get() ([]PoolMeta, error)
}

type CmdlinePoolMetaProvider struct {
	rmsUrl string
	pools  []string
}

func (s CmdlinePoolMetaProvider) Get() ([]PoolMeta, error) {
	ret := make([]PoolMeta, 0, len(s.pools)+1)
	for i, poolName := range s.pools {
		ret = append(ret, PoolMeta{alias: poolName, id: commM.ID(i)})
	}
	ret = append(ret, PoolMeta{alias: "", id: 0})
	return ret, nil
}

type tcmsResourcePoolProvider interface {
	ListResourcePools() ([]*commM.ResourcePool, error)
}

type TcmsPoolMetaProvider struct {
	tcmsClient tcmsResourcePoolProvider
}

func (s TcmsPoolMetaProvider) Get() ([]PoolMeta, error) {
	ret := make([]PoolMeta, 0)
	pools, err := s.tcmsClient.ListResourcePools()
	if err != nil {
		return nil, err
	}
	aliases := make(map[string]struct{}, len(pools))
	for _, pool := range pools {
		if pool.Aliases == "" {
			pool.Aliases = pool.Name
		}
		for _, alias := range strings.Split(pool.Aliases, ",") {
			alias = strings.TrimSpace(alias)
			if _, ok := aliases[alias]; ok {
				return nil, fmt.Errorf("duplicated alias: %s", alias)
			}
			aliases[alias] = struct{}{}
			ret = append(ret, PoolMeta{
				id:             pool.ID,
				name:           pool.Name,
				alias:          alias,
				rmsUrl:         pool.RmsURL,
				borrowedPolicy: pool.BorrowedPolicy,
			})
		}
	}
	return ret, nil
}

type ResourcePoolMgr struct {
	pools             []ResourcePoolInterface
	nonAvailablePools map[ResourcePoolInterface]struct{}
	aliasDict         map[string]commM.ID
	identifierDict    map[commM.ID]int
	poolMetaProvider  PoolMetaProvider
}

func (p *ResourcePoolMgr) init() {
	if p.pools == nil {
		p.pools = []ResourcePoolInterface{}
	}
	if p.nonAvailablePools == nil {
		p.nonAvailablePools = map[ResourcePoolInterface]struct{}{}
	}
	if p.aliasDict == nil {
		p.aliasDict = map[string]commM.ID{}
	}
	if p.identifierDict == nil {
		p.identifierDict = map[commM.ID]int{}
	}
}

func (p ResourcePoolMgr) String() string {
	var nonAvailablePools []string
	for pool := range p.nonAvailablePools {
		nonAvailablePools = append(nonAvailablePools, fmt.Sprintf("%s?pool=%s", pool.RMS().Url(), pool.Name()))
	}
	var pools []json.RawMessage
	for _, pool := range p.pools {
		pools = append(pools, json.RawMessage(pool.String()))
	}
	AliasDict := make(map[string]int)
	for alias, id := range p.aliasDict {
		if index, ok := p.identifierDict[id]; ok {
			AliasDict[alias] = index
		} else {
			AliasDict[alias] = -1
		}
	}
	printer := struct {
		Pools             []json.RawMessage `json:"pools"`
		NonAvailablePools []string          `json:"nonAvailablePools"`
		AliasDict         map[string]int    `json:"aliasDict"`
	}{
		Pools:             pools,
		NonAvailablePools: nonAvailablePools,
		AliasDict:         AliasDict,
	}
	marshal, err := json.Marshal(printer)
	if err != nil {
		return "err: " + err.Error()
	}
	return string(marshal)
}

func (p ResourcePoolMgr) GetResourcePoolByAlias(alias string) ResourcePoolInterface {
	if poolId, ok := p.aliasDict[alias]; ok {
		return p.GetResourcePoolById(poolId)
	}
	return nil
}

func (p ResourcePoolMgr) GetResourcePoolById(id commM.ID) ResourcePoolInterface {
	if index, ok := p.identifierDict[id]; ok {
		pool := p.pools[index]
		if _, ok := p.nonAvailablePools[pool]; ok {
			return nil
		}
		return pool
	}
	return nil
}

func (p ResourcePoolMgr) Copy() ResourcePoolMgr {
	t := p
	t.pools = make([]ResourcePoolInterface, len(p.pools))
	for i, pool := range p.pools {
		t.pools[i] = pool.Copy()
	}
	return t
}

func (p *ResourcePoolMgr) Refresh(ctx context.Context) error {
	err := p.refreshPoolMeta()
	if err != nil {
		return err
	}
	return p.refreshEveryPool(ctx)
}

func (p *ResourcePoolMgr) refreshPoolMeta() error {
	poolsMeta, err := p.poolMetaProvider.Get()
	if err != nil {
		return err
	}
	p.Merge(poolsMeta)
	return nil
}

func (p *ResourcePoolMgr) refreshEveryPool(ctx context.Context) error {
	p.nonAvailablePools = map[ResourcePoolInterface]struct{}{}
	for _, pool := range p.pools {
		err := pool.Refresh(ctx)
		if err != nil {
			klog.ErrorS(err, "refresh pool failed", "name", pool.Name())
			p.nonAvailablePools[pool] = struct{}{}
		}
	}
	return nil
}

func (p *ResourcePoolMgr) Merge(poolsMeta []PoolMeta) {
	p.addPools(poolsMeta)

	survivors := map[commM.ID]struct{}{}
	for _, meta := range poolsMeta {
		survivors[meta.id] = struct{}{}
	}
	p.removeUnusedPools(survivors)
}

func (p *ResourcePoolMgr) addPools(poolsMeta []PoolMeta) {
	for _, poolMeta := range poolsMeta {
		p.aliasDict[poolMeta.alias] = poolMeta.id
		if _, ok := p.identifierDict[poolMeta.id]; !ok {
			pool := &ResourcePool{
				id:     poolMeta.id,
				name:   poolMeta.name,
				rms:    rms{url: poolMeta.rmsUrl, httpClient: http.DefaultClient},
				assume: defaultAssume{store: JobStore},
			}
			pool.init()
			p.pools = append(p.pools, pool)
			index := len(p.pools) - 1
			p.identifierDict[poolMeta.id] = index
		}
	}
}

func (p *ResourcePoolMgr) removeUnusedPools(survivors map[commM.ID]struct{}) {
	newPools := make([]ResourcePoolInterface, 0, len(p.pools))
	for _, pool := range p.pools {
		if _, ok := survivors[pool.ID()]; ok {
			newPools = append(newPools, pool)
		}
	}
	newIdDict := map[commM.ID]int{}
	for i, pool := range newPools {
		newIdDict[pool.ID()] = i
	}
	newAliasDict := map[string]commM.ID{}
	for alias, id := range p.aliasDict {
		if _, ok := survivors[id]; ok {
			newAliasDict[alias] = p.aliasDict[alias]
		}
	}
	p.pools = newPools
	p.identifierDict = newIdDict
	p.aliasDict = newAliasDict
}

// Set respool to plan-exec and testbed depending on the testbed.ResourcePoolAlias, then return it.
func (p ResourcePoolMgr) SetRespool4PlanExec(exec *commM.TestPlanExecution, testbed *Testbed) (pool ResourcePoolInterface, err error) {
	if exec.ResourcepoolID != 0 {
		// has been set already
		pool = p.GetResourcePoolById(exec.ResourcepoolID)
	} else {
		pool = p.GetResourcePoolByAlias(testbed.ResourcePoolAlias())
	}
	if pool == nil {
		return nil, fmt.Errorf("can not find available resource pool for: repoolID #%d or testbed alias '%s'", exec.ResourcepoolID, testbed.ResourcePoolAlias())
	}

	testbed.SetResourcePool(pool)
	exec.ResourcepoolID = pool.ID()
	return pool, nil
}

type ResCheckStatus struct {
	Status  int
	Message string
}

var ResCheckEnough = ResCheckStatus{Status: ResCheckStatusEnough}

const (
	ResCheckStatusUnknown int = iota
	ResCheckStatusInsufficient
	ResCheckStatusEnough
	ResCheckStatusError
	ResCheckStatusTempDown
)

func (p ResourcePoolMgr) CheckResourceEnough(ctx JobContext, testbeds []Testbed) ResCheckStatus {
	// If multiple testbeds are in the same pool
	// The former testbed should be assumed to check the next testbed
	// CheckResourceEnough should not change anything
	// So if there is assuming, pool should be copied
	switch len(testbeds) {
	case 0:
		return ResCheckEnough
	case 1:
		// if there is only one testbed, assume can be skipped for performance
		return p.isResourceEnoughForSingleTestbed(testbeds[0])
	}

	t := p.Copy() // for temporary assume if multiple testbeds in the same pool
	for _, testbed := range testbeds {
		pool := t.GetResourcePoolByAlias(testbed.ResourcePoolAlias())
		if pool == nil {
			return ResCheckStatus{
				Status:  ResCheckStatusError,
				Message: fmt.Sprintf("can not find available resource pool for: %s", testbed.ResourcePoolAlias()),
			}
		}
		if t := pool.CheckResourceEnough(testbed.Request()); t != ResCheckEnough {
			return t
		}
		err := pool.Assume(ctx, testbed.FullName(), testbed.Request())
		if err != nil {
			klog.ErrorS(err, "CheckResourceEnough assume failed")
			return ResCheckStatus{
				Status:  ResCheckStatusError,
				Message: fmt.Sprintf("CheckResourceEnough assume failed: %s", err.Error()),
			}
		}
	}
	return ResCheckEnough
}

func (p ResourcePoolMgr) isResourceEnoughForSingleTestbed(testbed Testbed) ResCheckStatus {
	pool := p.GetResourcePoolByAlias(testbed.ResourcePoolAlias())
	if pool == nil {
		return ResCheckStatus{
			Status:  ResCheckStatusError,
			Message: fmt.Sprintf("can not find available resource pool for: %s", testbed.ResourcePoolAlias()),
		}
	}
	return pool.CheckResourceEnough(testbed.Request())
}

func (p ResourcePoolMgr) Assume(ctx JobContext, testbeds []Testbed) error {
	for _, testbed := range testbeds {
		if testbed.FullName() == "" {
			return fmt.Errorf("no testbed name for scheduling")
		}
		pool := p.GetResourcePoolByAlias(testbed.ResourcePoolAlias())
		if pool == nil {
			return fmt.Errorf("can not find this available resource pool: %s", testbed.ResourcePoolAlias())
		}
		err := pool.Assume(ctx, testbed.FullName(), testbed.Request())
		if err != nil {
			return err
		}
	}
	return nil
}

type ResourcePoolInterface interface {
	fmt.Stringer
	Refresh(ctx context.Context) error
	CheckResourceEnough(request PhysicalResource) ResCheckStatus
	Assume(ctx JobContext, testbedName string, request PhysicalResource) error
	Copy() ResourcePoolInterface
	ID() commM.ID
	Name() string
	RMS() RmsInterface
}

type ResourcePool struct {
	id       commM.ID
	name     string
	tempDown bool
	cache    ResourcePoolMetric
	rms      RmsInterface
	assume   AssumeInterface
}

func (p *ResourcePool) Name() string {
	return p.name
}

func (p *ResourcePool) RMS() RmsInterface {
	return p.rms
}

func (p *ResourcePool) init() {
	s := &p.cache
	if s.Testbeds == nil {
		s.Testbeds = map[string]PhysicalResource{}
	}
	if s.Assumed == nil {
		s.Assumed = map[string]AssumeRecord{}
	}
}

func (p ResourcePool) String() string {
	printer := struct {
		ID    commM.ID        `json:"id"`
		Name  string          `json:"name"`
		Cache json.RawMessage `json:"cache"`
	}{
		ID:    p.id,
		Name:  p.name,
		Cache: json.RawMessage(p.cache.String()),
	}
	marshal, err := json.Marshal(printer)
	if err != nil {
		return "err: " + err.Error()
	}
	return string(marshal)
}

func (p *ResourcePool) CheckResourceEnough(request PhysicalResource) ResCheckStatus {
	if p.tempDown {
		return ResCheckStatus{ResCheckStatusTempDown, fmt.Sprintf("pool '%s' is temporary down", p.name)}
	}
	if p.cache.Capacity.IsZero() { // todo special treatment for zero resource
		return ResCheckStatus{ResCheckStatusInsufficient, fmt.Sprintf("pool '%s' capacity is zero", p.name)}
	}
	if p.cache.GetFree().Sub(request).NonNegative() {
		return ResCheckEnough
	} else if !p.cache.Capacity.Sub(request).NonNegative() {
		return ResCheckStatus{
			Status:  ResCheckStatusError,
			Message: fmt.Sprintf("request %s exceed pool '%s' capacity: %s", request.String(), p.name, p.cache.Capacity.String()),
		}
	} else {
		return ResCheckStatus{
			Status:  ResCheckStatusInsufficient,
			Message: fmt.Sprintf("insufficient resource, requested: %s, pool '%s' free: %s", request.String(), p.name, p.cache.GetFree().String()),
		}
	}
}

func (p *ResourcePool) Assume(ctx JobContext, testbedName string, request PhysicalResource) error {
	if _, ok := p.cache.Testbeds[testbedName]; ok {
		return fmt.Errorf("conflict testbed name: %s", testbedName)
	}
	// these maps must be not nil
	p.cache.Testbeds[testbedName] = request
	p.cache.Assumed[testbedName] = p.assume.Assume(ctx, testbedName)
	return nil
}

func (p ResourcePool) Copy() ResourcePoolInterface {
	t := p
	t.cache = t.cache.Copy()
	return &t
}

func (p *ResourcePool) ID() commM.ID {
	return p.id
}

func (p *ResourcePool) Refresh(ctx context.Context) error {
	p.tempDown = false
	remote, err := p.rms.PullResourcePoolMetric(p.name)
	if err != nil {
		p.tempDown = true
		klog.ErrorS(err, "refresh pool fail")
		return nil
	}
	p.Merge(ctx, remote)
	return nil
}

func (p *ResourcePool) Merge(ctx context.Context, remote ResourcePoolMetric) {
	s := &p.cache
	s.Capacity = remote.Capacity
	s.NonTestbedOccupation = remote.NonTestbedOccupation
	var expiredAssume []string
	for name, record := range s.Assumed {
		if !p.assume.ShouldKeepAssume(context.TODO(), record) {
			expiredAssume = append(expiredAssume, name)
		}
	}
	for _, name := range expiredAssume {
		delete(s.Assumed, name)
	}

	for testbedName, resource := range remote.Testbeds {
		if _, ok := s.Assumed[testbedName]; ok {
			s.Testbeds[testbedName] = resource.Max(s.Testbeds[testbedName])
		} else {
			s.Testbeds[testbedName] = resource
		}
	}

	var livingTestbeds = map[string]struct{}{}
	for name := range s.Assumed {
		livingTestbeds[name] = struct{}{}
	}
	for name := range remote.Testbeds {
		livingTestbeds[name] = struct{}{}
	}

	var deadTestbeds []string
	for name := range s.Testbeds {
		if _, ok := livingTestbeds[name]; !ok {
			deadTestbeds = append(deadTestbeds, name)
		}
	}
	for _, name := range deadTestbeds {
		delete(s.Testbeds, name)
	}
}

type ResourcePoolMetric struct {
	Capacity             PhysicalResource            `json:"capacity"`
	NonTestbedOccupation PhysicalResource            `json:"nonTestbedOccupation"`
	Testbeds             map[string]PhysicalResource `json:"testbeds"`
	Assumed              map[string]AssumeRecord     `json:"assumed"`
}

type AssumeRecord struct {
	CreateAt   time.Time `json:"createAt"`
	OwnerJobID commM.ID  `json:"ownerJobID"`
}

type AssumeInterface interface {
	Assume(ctx JobContext, testbedName string) AssumeRecord
	ShouldKeepAssume(ctx context.Context, record AssumeRecord) bool
}

const MinAssumeDuration = 10 * time.Minute

type defaultAssume struct {
	store common.ArgoJobStore
}

func (d defaultAssume) Assume(ctx JobContext, _ string) AssumeRecord {
	return AssumeRecord{CreateAt: time.Now(), OwnerJobID: ctx.GetJob().ID()}
}

func (d defaultAssume) ShouldKeepAssume(ctx context.Context, record AssumeRecord) bool {
	if record.CreateAt.Add(MinAssumeDuration).After(time.Now()) {
		return true
	}
	job, err := d.store.GetJob(ctx, record.OwnerJobID)
	if err != nil {
		klog.ErrorS(err, "pick job during assume")
		return false
	}
	if job.Status == common.ArgoTestJobRunning {
		return true
	}
	return false
}

func (s ResourcePoolMetric) String() string {
	printer := struct {
		ResourcePoolMetric
		Free PhysicalResource `json:"free"`
	}{
		ResourcePoolMetric: s,
		Free:               s.GetFree(),
	}
	marshal, err := json.Marshal(printer)
	if err != nil {
		return "err: " + err.Error()
	}
	return string(marshal)
}

func (s ResourcePoolMetric) GetFree() PhysicalResource {
	var current = s.Capacity
	current = current.Sub(s.NonTestbedOccupation)

	for _, tb := range s.Testbeds {
		current = current.Sub(tb)
	}
	return current
}

func (s ResourcePoolMetric) Copy() ResourcePoolMetric {
	t := s
	if s.Assumed != nil {
		t.Assumed = map[string]AssumeRecord{}
		for k, v := range s.Assumed {
			t.Assumed[k] = v
		}
	}
	if s.Testbeds != nil {
		t.Testbeds = map[string]PhysicalResource{}
		for k, v := range s.Testbeds {
			t.Testbeds[k] = v
		}
	}
	return t
}

type RmsInterface interface {
	PullResourcePoolMetric(rmsResourcePool string) (ResourcePoolMetric, error)
	Url() string
}

type rms struct {
	url        string
	httpClient HTTPClient
}

type HTTPClient interface {
	Do(req *http.Request) (*http.Response, error)
}

func (r rms) PullResourcePoolMetric(rmsResourcePool string) (result ResourcePoolMetric, err error) {
	result.Capacity, result.NonTestbedOccupation, err = r.pullResourcePoolStatus(rmsResourcePool)
	if err != nil {
		return
	}
	result.Testbeds, err = r.pullTestbedsStatus(rmsResourcePool)
	if err != nil {
		return
	}
	return
}

func (r rms) pullResourcePoolStatus(rmsResourcePool string) (capacity, nonTestbedOccupation PhysicalResource, err error) {
	poolStatusReq := rmsReq.GetResourcePoolStatusRequest{ResourcePool: rmsResourcePool}
	poolStatusResp := &rmsResp.GetResourcePoolStatusResponse{}
	err = r.request("api/v1/resource-pool/status", poolStatusReq, poolStatusResp)
	if err != nil {
		return
	}
	capacity = PhysicalResource(poolStatusResp.TotalResource)
	nonTestbedOccupation = PhysicalResource(poolStatusResp.NonTestbedOccupation)
	return
}

func (r rms) pullTestbedsStatus(rmsResourcePool string) (result map[string]PhysicalResource, err error) {
	testbedsReq := rmsReq.GetTestbedsResourceRequest{ResourcePool: rmsResourcePool}
	testbedsResp := &rmsResp.GetTestbedsResourceResponse{}
	err = r.request("api/v1/resource-pool/testbeds-resource", testbedsReq, testbedsResp)
	if err != nil {
		return
	}
	result = map[string]PhysicalResource{}
	for name, resource := range testbedsResp.TestbedsOccupation {
		result[name] = PhysicalResource(*resource)
	}
	return
}

func (r rms) request(path string, req interface{}, resp interface{}) (err error) {
	buffer := &bytes.Buffer{}
	err = json.NewEncoder(buffer).Encode(req)
	if err != nil {
		return
	}
	httpReq, err := http.NewRequest(http.MethodGet, r.url+"/"+path, buffer)
	httpReq.Header.Set("Content-Type", "application/json")
	if err != nil {
		return
	}
	httpResp, err := r.httpClient.Do(httpReq)
	if err != nil {
		return
	}
	defer httpResp.Body.Close()
	if httpResp.StatusCode != http.StatusOK {
		body, err := io.ReadAll(httpResp.Body)
		if err != nil {
			return fmt.Errorf("bad http status %d: %s, body err: %w", httpResp.StatusCode, httpReq.URL.String(), err)
		}
		return fmt.Errorf("bad http status %d: %s, body: %s", httpResp.StatusCode, httpReq.URL.String(), string(body))
	}
	err = json.NewDecoder(httpResp.Body).Decode(resp)
	if err != nil {
		return
	}
	return
}

func (r rms) Url() string {
	return r.url
}

var Tcms common.TCMS
var JobStore common.ArgoJobStore