package connector

import (
	"context"
	"fmt"

	"example/baton-junction/pkg/client"

	v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2"
	"github.com/conductorone/baton-sdk/pkg/annotations"
	"github.com/conductorone/baton-sdk/pkg/types/entitlement"
	"github.com/conductorone/baton-sdk/pkg/types/grant"
	"github.com/conductorone/baton-sdk/pkg/types/resource"
	"github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap"
	"go.uber.org/zap"
)

type roleBuilder struct {
	client *client.Client
}

func newRoleBuilder(c *client.Client) *roleBuilder {
	return &roleBuilder{client: c}
}

func (o *roleBuilder) ResourceType(_ context.Context) *v2.ResourceType {
	return roleResourceType
}

func (o *roleBuilder) List(
	ctx context.Context,
	_ *v2.ResourceId,
	attrs resource.SyncOpAttrs,
) ([]*v2.Resource, *resource.SyncOpResults, error) {
	l := ctxzap.Extract(ctx)

	bag, cursor, err := parsePageToken(attrs.PageToken.Token, &v2.ResourceId{ResourceType: roleResourceType.Id})
	if err != nil {
		return nil, nil, err
	}

	roles, nextCursor, err := o.client.GetRoles(ctx, cursor)
	if err != nil {
		return nil, nil, fmt.Errorf("baton-junction: failed to list roles: %w", err)
	}

	l.Debug("listed roles", zap.Int("count", len(roles)), zap.String("next_cursor", nextCursor))

	var resources []*v2.Resource
	for _, r := range roles {
		res, err := roleResource(&r)
		if err != nil {
			return nil, nil, err
		}
		resources = append(resources, res)
	}

	if err := bag.Next(nextCursor); err != nil {
		return nil, nil, fmt.Errorf("baton-junction: failed to advance pagination: %w", err)
	}

	nextPageToken, err := bag.Marshal()
	if err != nil {
		return nil, nil, fmt.Errorf("baton-junction: failed to marshal pagination bag: %w", err)
	}

	return resources, &resource.SyncOpResults{NextPageToken: nextPageToken}, nil
}

func (o *roleBuilder) Entitlements(
	_ context.Context,
	res *v2.Resource,
	_ resource.SyncOpAttrs,
) ([]*v2.Entitlement, *resource.SyncOpResults, error) {
	e := entitlement.NewAssignmentEntitlement(
		res,
		memberEntitlement,
		entitlement.WithGrantableTo(userResourceType, groupResourceType, roleResourceType),
		entitlement.WithDisplayName(fmt.Sprintf("%s Member", res.DisplayName)),
		entitlement.WithDescription(fmt.Sprintf("Assigned the %s role", res.DisplayName)),
	)
	return []*v2.Entitlement{e}, nil, nil
}

func (o *roleBuilder) Grants(
	ctx context.Context,
	res *v2.Resource,
	attrs resource.SyncOpAttrs,
) ([]*v2.Grant, *resource.SyncOpResults, error) {
	l := ctxzap.Extract(ctx)

	bag, cursor, err := parsePageToken(attrs.PageToken.Token, &v2.ResourceId{
		ResourceType: roleResourceType.Id,
		Resource:     res.Id.Resource,
	})
	if err != nil {
		return nil, nil, err
	}

	members, nextCursor, err := o.client.GetRoleMembers(ctx, res.Id.Resource, cursor)
	if err != nil {
		return nil, nil, fmt.Errorf("baton-junction: failed to list role members: %w", err)
	}

	l.Debug("listed role members",
		zap.String("role_id", res.Id.Resource),
		zap.Int("count", len(members)),
	)

	var grants []*v2.Grant
	for _, m := range members {
		g := grant.NewGrant(
			res,
			memberEntitlement,
			&v2.ResourceId{
				ResourceType: userResourceType.Id,
				Resource:     m.UserID,
			},
		)
		grants = append(grants, g)
	}

	roleGroups, err := o.client.GetRoleGroups(ctx, res.Id.Resource)
	if err != nil {
		return nil, nil, fmt.Errorf("baton-junction: failed to list role group assignments: %w", err)
	}

	for _, ga := range roleGroups {
		g := grant.NewGrant(
			res,
			memberEntitlement,
			&v2.ResourceId{
				ResourceType: groupResourceType.Id,
				Resource:     ga.GroupID,
			},
			grant.WithAnnotation(&v2.GrantExpandable{
				EntitlementIds: []string{fmt.Sprintf("group:%s:member", ga.GroupID)},
			}),
		)
		grants = append(grants, g)
	}

	l.Debug("listed role group assignments",
		zap.String("role_id", res.Id.Resource),
		zap.Int("group_count", len(roleGroups)),
	)

	roleRoles, err := o.client.GetRoleRoles(ctx, res.Id.Resource)
	if err != nil {
		return nil, nil, fmt.Errorf("baton-junction: failed to list role-to-role assignments: %w", err)
	}

	for _, ra := range roleRoles {
		g := grant.NewGrant(
			res,
			memberEntitlement,
			&v2.ResourceId{
				ResourceType: roleResourceType.Id,
				Resource:     ra.RoleID,
			},
			grant.WithAnnotation(&v2.GrantExpandable{
				EntitlementIds: []string{fmt.Sprintf("role:%s:member", ra.RoleID)},
			}),
		)
		grants = append(grants, g)
	}

	l.Debug("listed role-to-role assignments",
		zap.String("role_id", res.Id.Resource),
		zap.Int("role_count", len(roleRoles)),
	)

	if err := bag.Next(nextCursor); err != nil {
		return nil, nil, fmt.Errorf("baton-junction: failed to advance pagination: %w", err)
	}

	nextPageToken, err := bag.Marshal()
	if err != nil {
		return nil, nil, fmt.Errorf("baton-junction: failed to marshal pagination bag: %w", err)
	}

	return grants, &resource.SyncOpResults{NextPageToken: nextPageToken}, nil
}

func (o *roleBuilder) Grant(
	ctx context.Context,
	principal *v2.Resource,
	en *v2.Entitlement,
) (annotations.Annotations, error) {
	l := ctxzap.Extract(ctx)

	roleID := en.Resource.Id.Resource
	userID := principal.Id.Resource

	l.Info("granting role membership",
		zap.String("role_id", roleID),
		zap.String("user_id", userID),
	)

	err := o.client.AddRoleMember(ctx, roleID, userID)
	if err != nil {
		return nil, fmt.Errorf("baton-junction: failed to grant role membership: %w", err)
	}

	return nil, nil
}

func (o *roleBuilder) Revoke(
	ctx context.Context,
	g *v2.Grant,
) (annotations.Annotations, error) {
	l := ctxzap.Extract(ctx)

	roleID := g.Entitlement.Resource.Id.Resource
	userID := g.Principal.Id.Resource

	l.Info("revoking role membership",
		zap.String("role_id", roleID),
		zap.String("user_id", userID),
	)

	if err := o.client.RemoveRoleMember(ctx, roleID, userID); err != nil {
		if isNotFound(err) {
			l.Debug("role membership already removed", zap.String("role_id", roleID), zap.String("user_id", userID))
			return nil, nil
		}
		return nil, fmt.Errorf("baton-junction: failed to revoke role membership: %w", err)
	}

	return nil, nil
}

func (o *roleBuilder) Get(
	ctx context.Context,
	resourceId *v2.ResourceId,
	_ *v2.ResourceId,
) (*v2.Resource, annotations.Annotations, error) {
	l := ctxzap.Extract(ctx)

	role, err := o.client.GetRole(ctx, resourceId.Resource)
	if err != nil {
		return nil, nil, fmt.Errorf("baton-junction: failed to get role %s: %w", resourceId.Resource, err)
	}

	l.Debug("fetched role", zap.String("role_id", resourceId.Resource))

	r, err := roleResource(role)
	if err != nil {
		return nil, nil, err
	}

	return r, nil, nil
}

func roleResource(r *client.Role) (*v2.Resource, error) {
	res, err := resource.NewRoleResource(
		r.Name,
		roleResourceType,
		r.ID,
		[]resource.RoleTraitOption{
			resource.WithRoleProfile(map[string]interface{}{
				"description": r.Description,
			}),
		},
		resource.WithAnnotation(&v2.RawId{Id: r.ID}),
		resource.WithDescription(r.Description),
	)
	if err != nil {
		return nil, fmt.Errorf("baton-junction: failed to build role resource %s: %w", r.ID, err)
	}

	return res, nil
}