package connection

import (
	"context"
	"fmt"
	"log"
	"time"

	pb "nexusrmm.local/agent/pkg/proto"
	"google.golang.org/grpc"
	"google.golang.org/grpc/connectivity"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/keepalive"
)

type GrpcClient struct {
	conn    *grpc.ClientConn
	Client  pb.AgentServiceClient
	address string
}

func NewGrpcClient(address string) (*GrpcClient, error) {
	opts := []grpc.DialOption{
		grpc.WithTransportCredentials(insecure.NewCredentials()),
		grpc.WithKeepaliveParams(keepalive.ClientParameters{
			Time:                30 * time.Second,
			Timeout:             10 * time.Second,
			PermitWithoutStream: true,
		}),
	}

	conn, err := grpc.NewClient(address, opts...)
	if err != nil {
		return nil, err
	}

	return &GrpcClient{
		conn:    conn,
		Client:  pb.NewAgentServiceClient(conn),
		address: address,
	}, nil
}

func (g *GrpcClient) Close() {
	if g.conn != nil {
		g.conn.Close()
	}
}

func ConnectWithRetry(address string, maxRetries int) (*GrpcClient, error) {
	for i := 0; i < maxRetries; i++ {
		client, err := NewGrpcClient(address)
		if err != nil {
			log.Printf("Client creation %d/%d failed: %v", i+1, maxRetries, err)
			time.Sleep(backoffDuration(i))
			continue
		}

		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		client.conn.Connect()
		state := client.conn.GetState()
		client.conn.WaitForStateChange(ctx, state)
		cancel()

		newState := client.conn.GetState()
		if newState == connectivity.Ready || newState == connectivity.Idle {
			return client, nil
		}

		log.Printf("Connection attempt %d/%d: state=%v", i+1, maxRetries, newState)
		client.Close()
		time.Sleep(backoffDuration(i))
	}
	return nil, fmt.Errorf("failed to connect after %d attempts", maxRetries)
}

func backoffDuration(attempt int) time.Duration {
	secs := 2 << attempt
	if secs > 30 {
		secs = 30
	}
	return time.Duration(secs) * time.Second
}