package snmp

import (
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"path/filepath"
	"strings"
	"time"

	g "github.com/gosnmp/gosnmp"
	"github.com/slayercat/GoSNMPServer"
	"github.com/slayercat/GoSNMPServer/mibImps"
)

type SNMPService struct {
	ListenAddr string
	ListenPort uint16
	UserName   string
	AuthPass   string
	AuthProto  string
	PrivPass   string
	PrivProto  string
	EngineID   string
	TrapPort   uint16
	TrapListen bool
	TrapBool   bool
	TrapTick   uint16
	TimeOut    uint16
	TrapTarget string

	ListenHost string
	TrapHost   string

	SysName     string
	SysDescr    string
	SysLocation string
	SysContact  string
	SysStatus   string
	SysService  int
}

func (s *SNMPService) getAuthProto() g.SnmpV3AuthProtocol {
	switch s.AuthProto {
	case "NoAuth":
		return g.NoAuth
	case "MD5":
		return g.MD5
	case "SHA":
		return g.SHA
	default:
	}
	return g.MD5
}

func (s *SNMPService) getPrivProto() g.SnmpV3PrivProtocol {
	switch s.PrivProto {
	case "NoPriv":
		return g.NoPriv
	case "DES":
		return g.DES
	case "AES":
		return g.AES
	case "AES192":
		return g.AES192
	case "AES256":
		return g.AES256
	default:
	}
	return g.DES
}

func (s *SNMPService) setSecParamsList() []g.UsmSecurityParameters {
	var secParamsList = []g.UsmSecurityParameters{
		{
			UserName:                 s.UserName,
			AuthenticationProtocol:   s.getAuthProto(),
			AuthenticationPassphrase: s.AuthPass,
			PrivacyProtocol:          s.getPrivProto(),
			PrivacyPassphrase:        s.PrivPass,
			AuthoritativeEngineID:    s.EngineID,
		},
		// {
		// 	UserName:                 "myuser2",
		// 	AuthenticationProtocol:   g.SHA,
		// 	AuthenticationPassphrase: "mypassword2",
		// 	PrivacyProtocol:          g.DES,
		// 	PrivacyPassphrase:        "myprivacy2",
		// 	AuthoritativeEngineID:    s.EngineID,
		// },
		// {
		// 	UserName:                 "myuser2",
		// 	AuthenticationProtocol:   g.MD5,
		// 	AuthenticationPassphrase: "mypassword2",
		// 	PrivacyProtocol:          g.AES,
		// 	PrivacyPassphrase:        "myprivacy2",
		// 	AuthoritativeEngineID:    s.EngineID,
		// },
	}
	return secParamsList
}

func (s *SNMPService) StartSNMPServer() {
	// 设置引擎启动次数和引varvar
	var engineBoots uint32 = 1
	//var engineTime uint32 = uint32(time.Now().Unix() % 2147483647) // 使用当前时间初始化
	//var engineTime uint32 = 3600 // 使用当前时间初始化
	master := GoSNMPServer.MasterAgent{
		Logger: GoSNMPServer.NewDefaultLogger(),
		SecurityConfig: GoSNMPServer.SecurityConfig{
			NoSecurity:               true,
			AuthoritativeEngineBoots: engineBoots,
			// OnGetAuthoritativeEngineTime: func() uint32 {
			// 	return engineTime
			// },
			//AuthoritativeEngineID:    GoSNMPServer.SNMPEngineID{EngineIDData: "0x800007DB03360102101100"},

			Users: s.setSecParamsList(),
		},
		SubAgents: []*GoSNMPServer.SubAgent{
			{
				UserErrorMarkPacket: false,
				CommunityIDs:        []string{"public", "private"}, // SNMPv1 and SNMPv2c community strings
				OIDs:                s.handleOIDs(),
				//OIDs:                mibImps.All(),
			},
		},
	}

	server := GoSNMPServer.NewSNMPServer(master)
	err := server.ListenUDP("udp", s.ListenHost)
	if err != nil {
		log.Fatalf("Error in listen: %+v", err)
	}
	server.ServeForever()
}

func (s *SNMPService) handleOIDs() []*GoSNMPServer.PDUValueControlItem {
	customOIDs := []*GoSNMPServer.PDUValueControlItem{
		{
			OID:  "1.3.6.1.4.1.1379.2.3.3.3.1.1.1.0",
			Type: g.OctetString,
			OnGet: func() (value interface{}, err error) {
				return s.SysName, nil
			},
			OnSet: func(value interface{}) error {
				// 将[]uint8转换为string
				if v, ok := value.([]uint8); ok {
					s.SysName = string(v)
					log.Printf("Set request for OID 1.3.6.1.4.1.1379.2.3.3.3.1.1.1.0 with value %v", s.SysName)
					return nil
				}
				return nil
			},
		},
		{
			OID:  ".1.3.6.1.4.1.1379.2.3.3.3.1.1.2.0",
			Type: g.OctetString,
			OnGet: func() (value interface{}, err error) {
				return s.SysStatus, nil
			},
			OnSet: func(value interface{}) error {
				// 将[]uint8转换为string
				if v, ok := value.([]uint8); ok {
					s.SysStatus = string(v)
					log.Printf("Set request for OID 1.3.6.1.4.1.1379.2.3.3.3.1.1.2.0 with value %v", s.SysStatus)
					return nil
				}
				return nil
			},
		},
		{
			OID:  ".1.3.6.1.4.1.1379.2.3.3.3.1.1.3.0",
			Type: g.OctetString,
			OnGet: func() (value interface{}, err error) {
				return s.SysDescr, nil
			},
			OnSet: func(value interface{}) error {
				// 将[]uint8转换为string
				if v, ok := value.([]uint8); ok {
					s.SysDescr = string(v)
					log.Printf("Set request for OID .1.3.6.1.4.1.1379.2.3.3.3.1.1.3.0 with value %v", s.SysDescr)
					return nil
				}
				return nil
			},
		},
		{
			OID:  ".1.3.6.1.4.1.1379.2.3.3.3.1.1.4.0",
			Type: g.OctetString,
			OnGet: func() (value interface{}, err error) {
				return s.SysLocation, nil
			},
			OnSet: func(value interface{}) error {
				// 将[]uint8转换为string
				if v, ok := value.([]uint8); ok {
					s.SysLocation = string(v)
					log.Printf("Set request for OID .1.3.6.1.4.1.1379.2.3.3.3.1.1.4.0 with value %v", s.SysLocation)
					return nil
				}
				return nil
			},
		},
		{
			OID:  ".1.3.6.1.4.1.1379.2.3.3.3.1.1.5.0",
			Type: g.OctetString,
			OnGet: func() (value interface{}, err error) {
				return s.SysContact, nil
			},
			OnSet: func(value interface{}) error {
				// 将[]uint8转换为string
				if v, ok := value.([]uint8); ok {
					s.SysContact = string(v)
					log.Printf("Set request for OID .1.3.6.1.4.1.1379.2.3.3.3.1.1.5.0 with value %v", s.SysContact)
					return nil
				}
				return nil
			},
		},
		{
			OID:  ".1.3.6.1.4.1.1379.2.3.3.3.1.1.7.0",
			Type: g.TimeTicks,
			OnGet: func() (value interface{}, err error) {
				return uint32(time.Now().Unix()), nil
			},
		},
		{
			OID:  ".1.3.6.1.4.1.1379.2.3.3.3.1.1.9.0",
			Type: g.Integer,
			OnGet: func() (value interface{}, err error) {
				return s.SysService, nil
			},
			OnSet: func(value interface{}) error {
				if v, ok := value.(int); ok {
					s.SysService = v
					log.Printf("Set request for OID .1.3.6.1.4.1.1379.2.3.3.3.1.1.9.0 with value %v", s.SysService)
					return nil
				}
				return nil
			},
		},
	}
	// 为 GETBULK 新增处理 OIDs
	bulkOIDs := &GoSNMPServer.PDUValueControlItem{
		OID:  ".1.3.6.1.4.1.1379.2.3.3.3.1.1", // 这里是您想要支持 GETBULK 的 OID 前缀
		Type: g.OctetString,
		OnGet: func() (value interface{}, err error) {
			// 假设我们返回一百度值，您可以根据您的实现进行调整
			values := []interface{}{s.SysName, s.SysStatus, s.SysDescr, s.SysLocation, s.SysContact, uint32(time.Now().Unix()), s.SysService} // 可以从其他结构中获取真实值
			return values, nil
		},
	}

	customOIDs = append(customOIDs, bulkOIDs)

	// 获取mibImps.All()返回的OID列表
	mibOIDs := mibImps.All()

	// 使用Map来检测并移除重复的OID
	oidMap := make(map[string]*GoSNMPServer.PDUValueControlItem)
	for _, oid := range customOIDs {
		oidMap[oid.OID] = oid
	}
	for _, oid := range mibOIDs {
		if _, exists := oidMap[oid.OID]; !exists {
			oidMap[oid.OID] = oid
		} else {
			log.Printf("Duplicate OID found: %s", oid.OID)
		}
	}

	// 将Map转换为Slice
	allOIDs := make([]*GoSNMPServer.PDUValueControlItem, 0, len(oidMap))
	for _, oid := range oidMap {
		allOIDs = append(allOIDs, oid)
	}

	return allOIDs
}

func (s *SNMPService) StartTrapServer() {
	flag.Usage = func() {
		fmt.Printf("Usage:\n")
		fmt.Printf("   %s\n", filepath.Base(os.Args[0]))
		flag.PrintDefaults()
	}

	tl := g.NewTrapListener()
	tl.OnNewTrap = s.MyTrapHandler

	usmTable := g.NewSnmpV3SecurityParametersTable(g.NewLogger(log.New(os.Stdout, "", 0)))
	for i := range s.setSecParamsList() {
		sp := &s.setSecParamsList()[i] // 使用指针
		err := usmTable.Add(sp.UserName, sp)
		if err != nil {
			usmTable.Logger.Print(err)
		}
	}

	// 设置引擎启动次数和引varvar
	//var engineBoots uint32 = 1
	// var engineTime uint32 = uint32(time.Now().Unix() % 2147483647) // 使用当前时间初始化
	//var engineTime uint32 = 3600 // 使用当前时间初始化
	gs := &g.GoSNMP{
		Target:        s.TrapTarget,
		Port:          s.TrapPort,
		Transport:     "udp",
		Timeout:       time.Duration(s.TimeOut) * time.Second, // 设置超时时间为x秒
		Version:       g.Version3,                             // Always using version3 for traps, only option that works with all SNMP versions simultaneously
		MsgFlags:      g.NoAuthNoPriv,
		SecurityModel: g.UserSecurityModel,
		SecurityParameters: &g.UsmSecurityParameters{
			UserName:                 s.UserName,
			AuthoritativeEngineID:    s.EngineID,
			AuthoritativeEngineBoots: 1,
			//AuthoritativeEngineTime:  3600,
			AuthenticationProtocol:   s.getAuthProto(),
			AuthenticationPassphrase: s.AuthPass,
			PrivacyProtocol:          s.getPrivProto(),
			PrivacyPassphrase:        s.PrivPass,
		},
		//TrapSecurityParametersTable: usmTable,
		ContextEngineID: s.EngineID,
		ContextName:     "v3test",
	}
	tl.Params = gs
	tl.Params.Logger = g.NewLogger(log.New(os.Stdout, "", 0))

	// 定时发送Trap
	if s.TrapBool {
		go s.SendPeriodicTraps(gs)
	}
	go s.monitorNetwork(gs)

	if s.TrapListen {
		err := tl.Listen(s.TrapHost)
		if err != nil {
			log.Panicf("error in listen: %s", err)
		}
	}
}

func (s *SNMPService) MyTrapHandler(packet *g.SnmpPacket, addr *net.UDPAddr) {
	log.Printf("got trapdata from %s\n", addr.IP)
	for _, v := range packet.Variables {
		switch v.Type {
		case g.OctetString:
			b := v.Value.([]byte)
			fmt.Printf("OID: %s, string: %x\n", v.Name, b)

		default:
			log.Printf("trap: %+v\n", v)
		}
	}
}

func (s *SNMPService) SendPeriodicTraps(gs *g.GoSNMP) {
	err := gs.Connect()
	if err != nil {
		log.Fatalf("Connect() err: %v", err)
	}
	defer gs.Conn.Close()

	ticker := time.NewTicker(time.Duration(s.TrapTick) * time.Second) // 每10秒发送一次Trap
	defer ticker.Stop()

	for range ticker.C { // 每x秒发送一次Trap
		trap := g.SnmpTrap{
			Variables: []g.SnmpPDU{
				{
					Name:  ".1.3.6.1.2.1.1.3.0",
					Type:  g.TimeTicks,
					Value: uint32(time.Now().Unix()),
				},
				{
					Name:  ".1.3.6.1.6.3.1.1.4.1.0",
					Type:  g.ObjectIdentifier,
					Value: ".1.3.6.1.6.3.1.1.5.1",
				},
			},
		}
		_, err = gs.SendTrap(trap)
		if err != nil {
			log.Printf("error sending trap: %s", err)
		} else {
			log.Printf("trap sent successfully")
		}
	}
}

// 1. 设备链路连接失败时发送Trap (LinkDown)
func (s *SNMPService) sendLinkDownTrap(gs *g.GoSNMP, ifIndex int, ifDescr string) {
	err := gs.Connect()
	if err != nil {
		log.Fatalf("Connect() err: %v", err)
	}
	defer gs.Conn.Close()

	trap := g.SnmpTrap{
		Variables: []g.SnmpPDU{
			{
				Name:  ".1.3.6.1.4.1.1379.2.3.3.3.3.1", // linkDown
				Type:  g.OctetString,
				Value: ".1.3.6.1.4.1.1379.2.3.3.3.3.1",
			},
			{
				Name:  ".1.3.6.1.4.1.1379.2.3.3.3.3.1.1", // ifIndex
				Type:  g.Integer,
				Value: ifIndex,
			},
			{
				Name:  ".1.3.6.1.4.1.1379.2.3.3.3.3.1.2", // ifDescr
				Type:  g.OctetString,
				Value: ifDescr,
			},
			{
				Name:  ".1.3.6.1.4.1.1379.2.3.3.3.4", // severity OID
				Type:  g.Integer,
				Value: 2, // event
			},
		},
	}

	_, err = gs.SendTrap(trap)
	if err != nil {
		log.Printf("error sending LinkDown trap: %s", err)
	} else {
		log.Printf("LinkDown trap sent successfully")
	}
}

// 2. 设备链路恢复正常时发送Trap (LinkUp)
func (s *SNMPService) sendLinkUpTrap(gs *g.GoSNMP, ifIndex int, ifDescr string) {
	err := gs.Connect()
	if err != nil {
		log.Fatalf("Connect() err: %v", err)
	}
	defer gs.Conn.Close()

	trap := g.SnmpTrap{
		Variables: []g.SnmpPDU{
			{
				Name:  ".1.3.6.1.4.1.1379.2.3.3.3.3.2", // linkUp
				Type:  g.OctetString,
				Value: ".1.3.6.1.4.1.1379.2.3.3.3.3.2",
			},
			{
				Name:  ".1.3.6.1.4.1.1379.2.3.3.3.3.2.1", // ifIndex
				Type:  g.Integer,
				Value: ifIndex,
			},
			{
				Name:  ".1.3.6.1.4.1.1379.2.3.3.3.3.2.2", // ifDescr
				Type:  g.OctetString,
				Value: ifDescr,
			},
			{
				Name:  ".1.3.6.1.4.1.1379.2.3.3.3.4", // severity OID
				Type:  g.Integer,
				Value: 5, // event
			},
		},
	}

	_, err = gs.SendTrap(trap)
	if err != nil {
		log.Printf("error sending LinkUp trap: %s", err)
	} else {
		log.Printf("LinkUp trap sent successfully")
	}
}

// 3. 设备鉴权失败时发送Trap (AuthenticationFailure)
func (s *SNMPService) sendAuthFailureTrap(gs *g.GoSNMP, username, descr string) {
	err := gs.Connect()
	if err != nil {
		log.Fatalf("Connect() err: %v", err)
	}
	defer gs.Conn.Close()

	trap := g.SnmpTrap{
		Variables: []g.SnmpPDU{
			{
				Name:  ".1.3.6.1.4.1.1379.2.3.3.3.3.3", // authenticationFailure
				Type:  g.OctetString,
				Value: ".1.3.6.1.4.1.1379.2.3.3.3.3.3",
			},
			{
				Name:  ".1.3.6.1.4.1.1379.2.3.3.3.3.3.1", // 自定义OID，用于记录失败的用户名
				Type:  g.OctetString,
				Value: username,
			},
			{
				Name:  ".1.3.6.1.4.1.1379.2.3.3.3.3.3.2", // 自定义OID，用于记录描述
				Type:  g.OctetString,
				Value: descr,
			},
			{
				Name:  ".1.3.6.1.4.1.1379.2.3.3.3.4", // severity OID
				Type:  g.Integer,
				Value: 4, // event
			},
		},
	}

	_, err = gs.SendTrap(trap)
	if err != nil {
		log.Printf("error sending AuthenticationFailure trap: %s", err)
	} else {
		log.Printf("AuthenticationFailure trap sent successfully")
	}
}

func (s *SNMPService) monitorNetwork(gs *g.GoSNMP) {
	// 假设有一个函数 checkLinkStatus 返回链路状态
	for {
		serviceStatus := s.checkServiceStatus()
		switch strings.ToUpper(serviceStatus) {
		case "LINK_DOWN":
			index := 1
			ifDescr := fmt.Sprintf("Link(index=%d) DOWN", index)
			s.sendLinkDownTrap(gs, index, ifDescr) // 假设接口索引为1
			s.SysService = 0
		case "LINK_UP":
			index := 1
			ifDescr := fmt.Sprintf("Link(index=%d) UP", index)
			s.sendLinkUpTrap(gs, index, ifDescr) // 假设接口索引为1
			s.SysService = 0
		case "AUTH_FAILURE":
			descr := "Authentication Failure"
			s.sendAuthFailureTrap(gs, s.UserName, descr)
			s.SysService = 0
		default:
		}

		time.Sleep(10 * time.Second) // 每10秒检查一次
	}
}

func (s *SNMPService) checkServiceStatus() string {
	switch s.SysService {
	case 1:
		return "LINK_DOWN"
	case 2:
		return "LINK_UP"
	case 3:
		return "AUTH_FAILURE"
	default:
	}
	return "NORMAL"
}