package mongodb

import (
	"context"
	"fmt"
	"io"
	"net"
	"strings"
	"time"

	"github.com/shadow1ng/fscan/common"
	"github.com/shadow1ng/fscan/plugins/base"
)

// MongoDBConnection MongoDB连接包装器
type MongoDBConnection struct {
	client net.Conn
	target string
}

// MongoDBConnector MongoDB连接器实现
type MongoDBConnector struct {
	host string
	port string
}

// NewMongoDBConnector 创建MongoDB连接器
func NewMongoDBConnector() *MongoDBConnector {
	return &MongoDBConnector{}
}

// Connect 连接到MongoDB服务
func (c *MongoDBConnector) Connect(ctx context.Context, info *common.HostInfo) (interface{}, error) {
	c.host = info.Host
	c.port = info.Ports
	
	target := fmt.Sprintf("%s:%s", c.host, c.port)
	timeout := time.Duration(common.Timeout) * time.Second
	
	// 建立TCP连接
	conn, err := net.DialTimeout("tcp", target, timeout)
	if err != nil {
		return nil, fmt.Errorf("MongoDB连接失败: %v", err)
	}
	
	return &MongoDBConnection{
		client: conn,
		target: target,
	}, nil
}

// Authenticate 认证 - MongoDB通常检查未授权访问
func (c *MongoDBConnector) Authenticate(ctx context.Context, conn interface{}, cred *base.Credential) error {
	mongoConn, ok := conn.(*MongoDBConnection)
	if !ok {
		return fmt.Errorf("无效的连接类型")
	}
	
	// 在goroutine中进行操作，支持Context取消
	resultChan := make(chan error, 1)
	
	go func() {
		// 设置操作超时
		timeout := time.Duration(common.Timeout) * time.Second
		if err := mongoConn.client.SetDeadline(time.Now().Add(timeout)); err != nil {
			resultChan <- fmt.Errorf("设置超时失败: %v", err)
			return
		}
		
		// 尝试未授权访问检测
		isUnauth, err := c.checkMongoDBUnauth(ctx, mongoConn)
		if err != nil {
			resultChan <- fmt.Errorf("MongoDB检测失败: %v", err)
			return
		}
		
		if isUnauth {
			// 未授权访问成功
			resultChan <- nil
			return
		}
		
		// 如果没有未授权访问，尝试使用提供的凭据（如果有的话）
		if cred.Username != "" && cred.Password != "" {
			// TODO: 实现MongoDB认证
			resultChan <- fmt.Errorf("MongoDB需要认证但暂不支持密码认证")
			return
		}
		
		resultChan <- fmt.Errorf("MongoDB服务需要认证")
	}()
	
	// 等待操作结果或Context取消
	select {
	case err := <-resultChan:
		return err
	case <-ctx.Done():
		return fmt.Errorf("MongoDB操作超时: %v", ctx.Err())
	}
}

// checkMongoDBUnauth 检测MongoDB未授权访问
func (c *MongoDBConnector) checkMongoDBUnauth(ctx context.Context, mongoConn *MongoDBConnection) (bool, error) {
	// 先尝试OP_MSG查询
	msgPacket := createOpMsgPacket()
	reply, err := c.sendMongoQuery(ctx, mongoConn, msgPacket)
	if err != nil {
		// 失败则尝试OP_QUERY查询
		queryPacket := createOpQueryPacket()
		reply, err = c.sendMongoQuery(ctx, mongoConn, queryPacket)
		if err != nil {
			return false, err
		}
	}
	
	// 检查响应结果
	if strings.Contains(reply, "totalLinesWritten") {
		return true, nil
	}
	
	return false, nil
}

// sendMongoQuery 发送MongoDB查询
func (c *MongoDBConnector) sendMongoQuery(ctx context.Context, mongoConn *MongoDBConnection, packet []byte) (string, error) {
	// 检查上下文是否已取消
	select {
	case <-ctx.Done():
		return "", ctx.Err()
	default:
	}
	
	// 发送查询包
	if _, err := mongoConn.client.Write(packet); err != nil {
		return "", fmt.Errorf("发送查询失败: %v", err)
	}
	
	// 再次检查上下文是否已取消
	select {
	case <-ctx.Done():
		return "", ctx.Err()
	default:
	}
	
	// 读取响应
	reply := make([]byte, 2048)
	count, err := mongoConn.client.Read(reply)
	if err != nil && err != io.EOF {
		return "", fmt.Errorf("读取响应失败: %v", err)
	}
	
	if count == 0 {
		return "", fmt.Errorf("收到空响应")
	}
	
	return string(reply[:count]), nil
}

// Close 关闭连接
func (c *MongoDBConnector) Close(conn interface{}) error {
	if mongoConn, ok := conn.(*MongoDBConnection); ok && mongoConn.client != nil {
		return mongoConn.client.Close()
	}
	return nil
}

// createOpMsgPacket 创建OP_MSG查询包
func createOpMsgPacket() []byte {
	return []byte{
		0x69, 0x00, 0x00, 0x00, // messageLength
		0x39, 0x00, 0x00, 0x00, // requestID
		0x00, 0x00, 0x00, 0x00, // responseTo
		0xdd, 0x07, 0x00, 0x00, // opCode OP_MSG
		0x00, 0x00, 0x00, 0x00, // flagBits
		// sections db.adminCommand({getLog: "startupWarnings"})
		0x00, 0x54, 0x00, 0x00, 0x00, 0x02, 0x67, 0x65, 0x74, 0x4c, 0x6f, 0x67, 0x00, 0x10, 0x00, 0x00, 0x00, 0x73, 0x74, 0x61, 0x72, 0x74, 0x75, 0x70, 0x57, 0x61, 0x72, 0x6e, 0x69, 0x6e, 0x67, 0x73, 0x00, 0x02, 0x24, 0x64, 0x62, 0x00, 0x06, 0x00, 0x00, 0x00, 0x61, 0x64, 0x6d, 0x69, 0x6e, 0x00, 0x03, 0x6c, 0x73, 0x69, 0x64, 0x00, 0x1e, 0x00, 0x00, 0x00, 0x05, 0x69, 0x64, 0x00, 0x10, 0x00, 0x00, 0x00, 0x04, 0x6e, 0x81, 0xf8, 0x8e, 0x37, 0x7b, 0x4c, 0x97, 0x84, 0x4e, 0x90, 0x62, 0x5a, 0x54, 0x3c, 0x93, 0x00, 0x00,
	}
}

// createOpQueryPacket 创建OP_QUERY查询包
func createOpQueryPacket() []byte {
	return []byte{
		0x48, 0x00, 0x00, 0x00, // messageLength
		0x02, 0x00, 0x00, 0x00, // requestID
		0x00, 0x00, 0x00, 0x00, // responseTo
		0xd4, 0x07, 0x00, 0x00, // opCode OP_QUERY
		0x00, 0x00, 0x00, 0x00, // flags
		0x61, 0x64, 0x6d, 0x69, 0x6e, 0x2e, 0x24, 0x63, 0x6d, 0x64, 0x00, // fullCollectionName admin.$cmd
		0x00, 0x00, 0x00, 0x00, // numberToSkip
		0x01, 0x00, 0x00, 0x00, // numberToReturn
		// query db.adminCommand({getLog: "startupWarnings"})
		0x21, 0x00, 0x00, 0x00, 0x2, 0x67, 0x65, 0x74, 0x4c, 0x6f, 0x67, 0x00, 0x10, 0x00, 0x00, 0x00, 0x73, 0x74, 0x61, 0x72, 0x74, 0x75, 0x70, 0x57, 0x61, 0x72, 0x6e, 0x69, 0x6e, 0x67, 0x73, 0x00, 0x00,
	}
}