package grpcclient

import (
	"context"
	"crypto/tls"
	"time"

	"github.com/mrhid6/keymanager/agent/internal/grpc/pb"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/encoding"
)

func init() {
	encoding.RegisterCodec(JSONCodec{})
}

type Client struct {
	conn   *grpc.ClientConn
	client pb.KeyManagerClient
}

func New(serverURL string, useTLS bool) (*Client, error) {
	var dialOpts []grpc.DialOption

	if useTLS {
		tlsCfg := &tls.Config{
			InsecureSkipVerify: false,
		}
		creds := credentials.NewTLS(tlsCfg)
		dialOpts = append(dialOpts, grpc.WithTransportCredentials(creds))
	} else {
		dialOpts = append(dialOpts, grpc.WithTransportCredentials(insecure.NewCredentials()))
	}

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	conn, err := grpc.DialContext(ctx, serverURL, dialOpts...)
	if err != nil {
		return nil, err
	}

	return &Client{
		conn:   conn,
		client: pb.NewKeyManagerClient(conn),
	}, nil
}

func (c *Client) Close() error {
	return c.conn.Close()
}

func (c *Client) Register(serverID, preRegToken, hostname, ipAddress, osInfo string) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	resp, err := c.client.Register(ctx, &pb.RegisterRequest{
		ServerId:    serverID,
		PreRegToken: preRegToken,
		Hostname:    hostname,
		IpAddress:   ipAddress,
		OsInfo:      osInfo,
	})
	if err != nil {
		return "", err
	}
	return resp.AgentToken, nil
}

func (c *Client) SyncKeys(serverID, agentToken string) ([]string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	resp, err := c.client.SyncKeys(ctx, &pb.SyncRequest{
		ServerId:   serverID,
		AgentToken: agentToken,
	})
	if err != nil {
		return nil, err
	}
	return resp.PublicKeys, nil
}

func (c *Client) UploadGeneratedKey(serverID, agentToken, publicKey, label string) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	resp, err := c.client.UploadGeneratedKey(ctx, &pb.UploadKeyRequest{
		ServerId:   serverID,
		AgentToken: agentToken,
		PublicKey:  publicKey,
		Label:      label,
	})
	if err != nil {
		return "", err
	}
	return resp.KeyId, nil
}