package parser

import (
	"errors"
	"fmt"
	"io"
	"net"

	"github.com/chainreactors/logs"
	"github.com/chainreactors/malice-network/helper/consts"
	"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/server/internal/parser/malefic"
	"github.com/chainreactors/malice-network/server/internal/parser/pulse"
)

// DetectProtocol 检测协议类型
func DetectProtocol(data []byte) (*MessageParser, error) {
	// Malefic 协议以 0xd1 开头
	if data[0] == malefic.DefaultStartDelimiter {
		return NewParser(consts.ImplantMalefic)
	}

	// Pulse 协议以 0x41 开头
	if data[0] == pulse.DefaultStartDelimiter {
		return NewParser(consts.ImplantPulse)
	}

	return nil, errors.New("unknown protocol")
}

// PacketParser packet parser, like malefic, beacon ...
type PacketParser interface {
	//PeekHeader(conn io.ReadWriteCloser) (uint32, uint32, error)
	ReadHeader(conn io.ReadWriteCloser) (uint32, uint32, error)
	Parse([]byte) (*implantpb.Spites, error)
	Marshal(*implantpb.Spites, uint32) ([]byte, error)
}

func NewParser(name string) (*MessageParser, error) {
	switch name {
	case consts.ImplantMalefic:
		return &MessageParser{Implant: name, PacketParser: malefic.NewMaleficParser()}, nil
	case consts.ImplantPulse:
		return &MessageParser{Implant: name, PacketParser: pulse.NewPulseParser()}, nil
	default:
		return nil, errs.ErrInvalidImplant
	}
}

type MessageParser struct {
	Implant string
	PacketParser
}

// WithSecure 为 MessageParser 添加安全支持
func (mp *MessageParser) WithSecure(keyPair *clientpb.KeyPair) {
	switch mp.Implant {
	case consts.ImplantMalefic:
		if maleficParser, ok := mp.PacketParser.(*malefic.MaleficParser); ok {
			maleficParser.WithSecure(keyPair)
		}
	default:

	}
}

func (parser *MessageParser) ReadMessage(conn io.ReadWriteCloser, length uint32) (*implantpb.Spites, error) {
	buf := make([]byte, length)
	_, err := io.ReadFull(conn, buf)
	if err != nil {
		return nil, err
	}
	return parser.Parse(buf)
}

func (parser *MessageParser) ReadPacket(conn io.ReadWriteCloser) (uint32, *implantpb.Spites, error) {
	sessionId, length, err := parser.ReadHeader(conn)
	if err != nil {
		return 0, nil, err
	}

	buf := make([]byte, length)
	_, err = io.ReadFull(conn, buf)
	if err != nil {
		return 0, nil, err
	}

	msg, err := parser.Parse(buf)
	return sessionId, msg, nil
}

func (parser *MessageParser) WritePacket(conn net.Conn, msg *implantpb.Spites, sid uint32) error {
	bs, err := parser.Marshal(msg, sid)
	if err != nil {
		return err
	}

	n, err := conn.Write(bs)
	if err != nil {
		return err
	}
	if n != len(bs) {
		return fmt.Errorf("send error, expect send %d, real send %d", len(bs), n)
	}
	if len(bs) <= 1000 {
		logs.Log.Debugf("write packet to %s , %d bytes", conn.RemoteAddr(), len(bs))
	} else {
		logs.Log.Debugf("write packet to %s , %d bytes", conn.RemoteAddr(), len(bs))
	}

	return nil
}
