probo/pkg/server/api/mcp/v1/schema.resolvers.go

package mcp_v1

// This file will be automatically regenerated based on the schema, any resolver implementations
// will be copied through when generating and any unknown code will be moved to the end.
// Code generated by mcpgen. DO NOT EDIT.

import (
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"time"

	"github.com/modelcontextprotocol/go-sdk/mcp"
	"go.probo.inc/probo/pkg/coredata"
	"go.probo.inc/probo/pkg/gid"
	"go.probo.inc/probo/pkg/iam"
	"go.probo.inc/probo/pkg/mail"
	"go.probo.inc/probo/pkg/page"
	"go.probo.inc/probo/pkg/probo"
	"go.probo.inc/probo/pkg/server/api/authn"
	"go.probo.inc/probo/pkg/server/api/mcp/v1/types"
)

// ListOrganizationsTool handles the listOrganizations tool
// List all organizations the user has access to
func (r *Resolver) ListOrganizationsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListOrganizationsInput) (*mcp.CallToolResult, types.ListOrganizationsOutput, error) {
	user := authn.IdentityFromContext(ctx)

	organizations, err := r.iamSvc.AccountService.ListOrganizations(ctx, user.ID)
	if err != nil {
		return nil, types.ListOrganizationsOutput{}, fmt.Errorf("failed to list organizations: %w", err)
	}

	result := types.ListOrganizationsOutput{
		Organizations: make([]*types.Organization, 0, len(organizations)),
	}

	for _, org := range organizations {
		result.Organizations = append(result.Organizations, types.NewOrganization(org))
	}

	return nil, result, nil
}

// ListVendorsTool handles the listVendors tool
// List all vendors for the organization
func (r *Resolver) ListVendorsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListVendorsInput) (*mcp.CallToolResult, types.ListVendorsOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionVendorList)

	prb := r.ProboService(ctx, input.OrganizationID)

	pageOrderBy := page.OrderBy[coredata.VendorOrderField]{
		Field:     coredata.VendorOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.VendorOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	noSnapshot := (*gid.GID)(nil)
	vendorFilter := coredata.NewVendorFilter(&noSnapshot, nil)
	if input.Filter != nil {
		vendorFilter = coredata.NewVendorFilter(&input.Filter.SnapshotID, nil)
	}

	page, err := prb.Vendors.ListForOrganizationID(ctx, input.OrganizationID, cursor, vendorFilter)
	if err != nil {
		panic(fmt.Errorf("cannot list organization vendors: %w", err))
	}

	return nil, types.NewListVendorsOutput(page), nil
}

// AddVendorTool handles the addVendor tool
// Add a new vendor to the organization
func (r *Resolver) AddVendorTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddVendorInput) (*mcp.CallToolResult, types.AddVendorOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionVendorCreate)

	svc := r.ProboService(ctx, input.OrganizationID)

	var category *coredata.VendorCategory
	if input.Category != nil {
		cat := coredata.VendorCategory(*input.Category)
		category = &cat
	}

	var countries coredata.CountryCodes
	if len(input.Countries) > 0 {
		countries = make(coredata.CountryCodes, len(input.Countries))
		for i, c := range input.Countries {
			countries[i] = coredata.CountryCode(c)
		}
	}

	vendor, err := svc.Vendors.Create(
		ctx,
		probo.CreateVendorRequest{
			OrganizationID:                input.OrganizationID,
			Name:                          input.Name,
			Description:                   input.Description,
			Category:                      category,
			HeadquarterAddress:            input.HeadquarterAddress,
			LegalName:                     input.LegalName,
			WebsiteURL:                    input.WebsiteURL,
			PrivacyPolicyURL:              input.PrivacyPolicyURL,
			ServiceLevelAgreementURL:      input.ServiceLevelAgreementURL,
			DataProcessingAgreementURL:    input.DataProcessingAgreementURL,
			BusinessAssociateAgreementURL: input.BusinessAssociateAgreementURL,
			SubprocessorsListURL:          input.SubprocessorsListURL,
			Certifications:                input.Certifications,
			Countries:                     countries,
			BusinessOwnerID:               input.BusinessOwnerID,
			SecurityOwnerID:               input.SecurityOwnerID,
			StatusPageURL:                 input.StatusPageURL,
			TermsOfServiceURL:             input.TermsOfServiceURL,
			SecurityPageURL:               input.SecurityPageURL,
			TrustPageURL:                  input.TrustPageURL,
		},
	)
	if err != nil {
		return nil, types.AddVendorOutput{}, fmt.Errorf("failed to create vendor: %w", err)
	}

	return nil, types.NewAddVendorOutput(vendor), nil
}

// UpdateVendorTool handles the updateVendor tool
// Update an existing vendor
func (r *Resolver) UpdateVendorTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateVendorInput) (*mcp.CallToolResult, types.UpdateVendorOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionVendorUpdate)

	svc := r.ProboService(ctx, input.ID)

	var description **string
	if input.Description != nil {
		description = &input.Description
	}

	var headquarterAddress **string
	if input.HeadquarterAddress != nil {
		headquarterAddress = &input.HeadquarterAddress
	}

	var legalName **string
	if input.LegalName != nil {
		legalName = &input.LegalName
	}

	var websiteURL **string
	if input.WebsiteURL != nil {
		websiteURL = &input.WebsiteURL
	}

	var privacyPolicyURL **string
	if input.PrivacyPolicyURL != nil {
		privacyPolicyURL = &input.PrivacyPolicyURL
	}

	var serviceLevelAgreementURL **string
	if input.ServiceLevelAgreementURL != nil {
		serviceLevelAgreementURL = &input.ServiceLevelAgreementURL
	}

	var dataProcessingAgreementURL **string
	if input.DataProcessingAgreementURL != nil {
		dataProcessingAgreementURL = &input.DataProcessingAgreementURL
	}

	var businessAssociateAgreementURL **string
	if input.BusinessAssociateAgreementURL != nil {
		businessAssociateAgreementURL = &input.BusinessAssociateAgreementURL
	}

	var subprocessorsListURL **string
	if input.SubprocessorsListURL != nil {
		subprocessorsListURL = &input.SubprocessorsListURL
	}

	var statusPageURL **string
	if input.StatusPageURL != nil {
		statusPageURL = &input.StatusPageURL
	}

	var termsOfServiceURL **string
	if input.TermsOfServiceURL != nil {
		termsOfServiceURL = &input.TermsOfServiceURL
	}

	var securityPageURL **string
	if input.SecurityPageURL != nil {
		securityPageURL = &input.SecurityPageURL
	}

	var trustPageURL **string
	if input.TrustPageURL != nil {
		trustPageURL = &input.TrustPageURL
	}

	var businessOwnerID **gid.GID
	if input.BusinessOwnerID != nil {
		businessOwnerID = &input.BusinessOwnerID
	}

	var securityOwnerID **gid.GID
	if input.SecurityOwnerID != nil {
		securityOwnerID = &input.SecurityOwnerID
	}

	var category *coredata.VendorCategory
	if input.Category != nil {
		cat := coredata.VendorCategory(*input.Category)
		category = &cat
	}

	var countries coredata.CountryCodes
	if len(input.Countries) > 0 {
		countries = make(coredata.CountryCodes, len(input.Countries))
		for i, c := range input.Countries {
			countries[i] = coredata.CountryCode(c)
		}
	}

	vendor, err := svc.Vendors.Update(
		ctx,
		probo.UpdateVendorRequest{
			ID:                            input.ID,
			Name:                          input.Name,
			Description:                   description,
			Category:                      category,
			HeadquarterAddress:            headquarterAddress,
			LegalName:                     legalName,
			WebsiteURL:                    websiteURL,
			PrivacyPolicyURL:              privacyPolicyURL,
			ServiceLevelAgreementURL:      serviceLevelAgreementURL,
			DataProcessingAgreementURL:    dataProcessingAgreementURL,
			BusinessAssociateAgreementURL: businessAssociateAgreementURL,
			SubprocessorsListURL:          subprocessorsListURL,
			Certifications:                input.Certifications,
			Countries:                     countries,
			BusinessOwnerID:               businessOwnerID,
			SecurityOwnerID:               securityOwnerID,
			StatusPageURL:                 statusPageURL,
			TermsOfServiceURL:             termsOfServiceURL,
			SecurityPageURL:               securityPageURL,
			TrustPageURL:                  trustPageURL,
		},
	)
	if err != nil {
		return nil, types.UpdateVendorOutput{}, fmt.Errorf("failed to update vendor: %w", err)
	}

	return nil, types.NewUpdateVendorOutput(vendor), nil
}

func (r *Resolver) ListRisksTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListRisksInput) (*mcp.CallToolResult, types.ListRisksOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionRiskList)

	prb := r.ProboService(ctx, input.OrganizationID)

	pageOrderBy := page.OrderBy[coredata.RiskOrderField]{
		Field:     coredata.RiskOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.RiskOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	noSnapshot := (*gid.GID)(nil)
	riskFilter := coredata.NewRiskFilter(nil, &noSnapshot)
	if input.Filter != nil {
		riskFilter = coredata.NewRiskFilter(input.Filter.Query, &input.Filter.SnapshotID)
	}

	page, err := prb.Risks.ListForOrganizationID(ctx, input.OrganizationID, cursor, riskFilter)
	if err != nil {
		panic(fmt.Errorf("cannot list organization risks: %w", err))
	}

	return nil, types.NewListRisksOutput(page), nil
}

func (r *Resolver) GetRiskTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetRiskInput) (*mcp.CallToolResult, types.GetRiskOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionRiskGet)

	prb := r.ProboService(ctx, input.ID)

	risk, err := prb.Risks.Get(ctx, input.ID)
	if err != nil {
		return nil, types.GetRiskOutput{}, fmt.Errorf("failed to get risk: %w", err)
	}

	return nil, types.GetRiskOutput{
		Risk: types.NewRisk(risk),
	}, nil
}

func (r *Resolver) AddRiskTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddRiskInput) (*mcp.CallToolResult, types.AddRiskOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionRiskCreate)

	svc := r.ProboService(ctx, input.OrganizationID)

	risk, err := svc.Risks.Create(
		ctx,
		probo.CreateRiskRequest{
			OrganizationID:     input.OrganizationID,
			Name:               input.Name,
			Description:        input.Description,
			Category:           input.Category,
			Treatment:          input.Treatment,
			InherentLikelihood: input.InherentLikelihood,
			InherentImpact:     input.InherentImpact,
			ResidualLikelihood: input.ResidualLikelihood,
			ResidualImpact:     input.ResidualImpact,
		},
	)
	if err != nil {
		return nil, types.AddRiskOutput{}, fmt.Errorf("failed to create risk: %w", err)
	}

	return nil, types.AddRiskOutput{
		Risk: types.NewRisk(risk),
	}, nil
}

func (r *Resolver) UpdateRiskTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateRiskInput) (*mcp.CallToolResult, types.UpdateRiskOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionRiskUpdate)

	svc := r.ProboService(ctx, input.ID)

	risk, err := svc.Risks.Update(
		ctx,
		probo.UpdateRiskRequest{
			ID:                 input.ID,
			Name:               input.Name,
			Description:        UnwrapOmittable(input.Description),
			Category:           input.Category,
			Treatment:          input.Treatment,
			OwnerID:            UnwrapOmittable(input.OwnerID),
			InherentLikelihood: input.InherentLikelihood,
			InherentImpact:     input.InherentImpact,
			ResidualLikelihood: input.ResidualLikelihood,
			ResidualImpact:     input.ResidualImpact,
		},
	)
	if err != nil {
		return nil, types.UpdateRiskOutput{}, fmt.Errorf("failed to update risk: %w", err)
	}

	return nil, types.UpdateRiskOutput{
		Risk: types.NewRisk(risk),
	}, nil
}

func (r *Resolver) ListMeasuresTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListMeasuresInput) (*mcp.CallToolResult, types.ListMeasuresOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionMeasureList)

	prb := r.ProboService(ctx, input.OrganizationID)

	pageOrderBy := page.OrderBy[coredata.MeasureOrderField]{
		Field:     coredata.MeasureOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.MeasureOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	var measureFilter = coredata.NewMeasureFilter(nil, nil)
	if input.Filter != nil {
		measureFilter = coredata.NewMeasureFilter(input.Filter.Query, input.Filter.State)
	}

	page, err := prb.Measures.ListForOrganizationID(ctx, input.OrganizationID, cursor, measureFilter)
	if err != nil {
		panic(fmt.Errorf("cannot list organization measures: %w", err))
	}

	return nil, types.NewListMeasuresOutput(page), nil
}

func (r *Resolver) GetMeasureTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetMeasureInput) (*mcp.CallToolResult, types.GetMeasureOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionMeasureGet)

	prb := r.ProboService(ctx, input.ID)

	measure, err := prb.Measures.Get(ctx, input.ID)
	if err != nil {
		return nil, types.GetMeasureOutput{}, fmt.Errorf("failed to get measure: %w", err)
	}

	return nil, types.GetMeasureOutput{
		Measure: types.NewMeasure(measure),
	}, nil
}

func (r *Resolver) AddMeasureTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddMeasureInput) (*mcp.CallToolResult, types.AddMeasureOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionMeasureCreate)

	svc := r.ProboService(ctx, input.OrganizationID)

	measure, err := svc.Measures.Create(
		ctx,
		probo.CreateMeasureRequest{
			OrganizationID: input.OrganizationID,
			Name:           input.Name,
			Description:    input.Description,
			Category:       input.Category,
		},
	)
	if err != nil {
		return nil, types.AddMeasureOutput{}, fmt.Errorf("failed to create measure: %w", err)
	}

	return nil, types.AddMeasureOutput{
		Measure: types.NewMeasure(measure),
	}, nil
}

func (r *Resolver) UpdateMeasureTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateMeasureInput) (*mcp.CallToolResult, types.UpdateMeasureOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionMeasureUpdate)

	svc := r.ProboService(ctx, input.ID)

	measure, err := svc.Measures.Update(
		ctx,
		probo.UpdateMeasureRequest{
			ID:          input.ID,
			Name:        input.Name,
			Description: UnwrapOmittable(input.Description),
			Category:    input.Category,
			State:       input.State,
		},
	)
	if err != nil {
		return nil, types.UpdateMeasureOutput{}, fmt.Errorf("failed to update measure: %w", err)
	}

	return nil, types.UpdateMeasureOutput{
		Measure: types.NewMeasure(measure),
	}, nil
}

func (r *Resolver) ListFrameworksTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListFrameworksInput) (*mcp.CallToolResult, types.ListFrameworksOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionFrameworkList)

	prb := r.ProboService(ctx, input.OrganizationID)

	pageOrderBy := page.OrderBy[coredata.FrameworkOrderField]{
		Field:     coredata.FrameworkOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.FrameworkOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	page, err := prb.Frameworks.ListForOrganizationID(ctx, input.OrganizationID, cursor)
	if err != nil {
		panic(fmt.Errorf("cannot list organization frameworks: %w", err))
	}

	return nil, types.NewListFrameworksOutput(page), nil
}

func (r *Resolver) GetFrameworkTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetFrameworkInput) (*mcp.CallToolResult, types.GetFrameworkOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionFrameworkGet)

	prb := r.ProboService(ctx, input.ID)

	framework, err := prb.Frameworks.Get(ctx, input.ID)
	if err != nil {
		return nil, types.GetFrameworkOutput{}, fmt.Errorf("failed to get framework: %w", err)
	}

	return nil, types.GetFrameworkOutput{
		Framework: types.NewFramework(framework),
	}, nil
}

func (r *Resolver) AddFrameworkTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddFrameworkInput) (*mcp.CallToolResult, types.AddFrameworkOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionFrameworkCreate)

	svc := r.ProboService(ctx, input.OrganizationID)

	framework, err := svc.Frameworks.Create(
		ctx,
		probo.CreateFrameworkRequest{
			OrganizationID: input.OrganizationID,
			Name:           input.Name,
			Description:    input.Description,
		},
	)
	if err != nil {
		return nil, types.AddFrameworkOutput{}, fmt.Errorf("failed to create framework: %w", err)
	}

	return nil, types.AddFrameworkOutput{
		Framework: types.NewFramework(framework),
	}, nil
}

func (r *Resolver) UpdateFrameworkTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateFrameworkInput) (*mcp.CallToolResult, types.UpdateFrameworkOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionFrameworkUpdate)

	svc := r.ProboService(ctx, input.ID)

	framework, err := svc.Frameworks.Update(
		ctx,
		probo.UpdateFrameworkRequest{
			ID:          input.ID,
			Name:        input.Name,
			Description: UnwrapOmittable(input.Description),
		},
	)
	if err != nil {
		return nil, types.UpdateFrameworkOutput{}, fmt.Errorf("failed to update framework: %w", err)
	}

	return nil, types.UpdateFrameworkOutput{
		Framework: types.NewFramework(framework),
	}, nil
}

func (r *Resolver) ListAssetsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListAssetsInput) (*mcp.CallToolResult, types.ListAssetsOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionAssetList)

	prb := r.ProboService(ctx, input.OrganizationID)

	pageOrderBy := page.OrderBy[coredata.AssetOrderField]{
		Field:     coredata.AssetOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.AssetOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	noSnapshot := (*gid.GID)(nil)
	assetFilter := coredata.NewAssetFilter(&noSnapshot)
	if input.Filter != nil {
		assetFilter = coredata.NewAssetFilter(&input.Filter.SnapshotID)
	}

	page, err := prb.Assets.ListForOrganizationID(ctx, input.OrganizationID, cursor, assetFilter)
	if err != nil {
		panic(fmt.Errorf("cannot list organization assets: %w", err))
	}

	return nil, types.NewListAssetsOutput(page), nil
}

func (r *Resolver) GetAssetTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetAssetInput) (*mcp.CallToolResult, types.GetAssetOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionAssetGet)

	prb := r.ProboService(ctx, input.ID)

	asset, err := prb.Assets.Get(ctx, input.ID)
	if err != nil {
		return nil, types.GetAssetOutput{}, fmt.Errorf("failed to get asset: %w", err)
	}

	return nil, types.GetAssetOutput{
		Asset: types.NewAsset(asset),
	}, nil
}

func (r *Resolver) AddAssetTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddAssetInput) (*mcp.CallToolResult, types.AddAssetOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionAssetCreate)

	svc := r.ProboService(ctx, input.OrganizationID)

	asset, err := svc.Assets.Create(
		ctx,
		probo.CreateAssetRequest{
			OrganizationID:  input.OrganizationID,
			Name:            input.Name,
			Amount:          input.Amount,
			OwnerID:         input.OwnerID,
			AssetType:       input.AssetType,
			DataTypesStored: input.DataTypesStored,
			VendorIDs:       input.VendorIds,
		},
	)
	if err != nil {
		return nil, types.AddAssetOutput{}, fmt.Errorf("failed to create asset: %w", err)
	}

	return nil, types.AddAssetOutput{
		Asset: types.NewAsset(asset),
	}, nil
}

func (r *Resolver) UpdateAssetTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateAssetInput) (*mcp.CallToolResult, types.UpdateAssetOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionAssetUpdate)

	svc := r.ProboService(ctx, input.ID)

	asset, err := svc.Assets.Update(
		ctx,
		probo.UpdateAssetRequest{
			ID:              input.ID,
			Name:            input.Name,
			Amount:          input.Amount,
			OwnerID:         input.OwnerID,
			AssetType:       input.AssetType,
			DataTypesStored: input.DataTypesStored,
			VendorIDs:       input.VendorIds,
		},
	)
	if err != nil {
		return nil, types.UpdateAssetOutput{}, fmt.Errorf("failed to update asset: %w", err)
	}

	return nil, types.UpdateAssetOutput{
		Asset: types.NewAsset(asset),
	}, nil
}

func (r *Resolver) ListDataTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListDataInput) (*mcp.CallToolResult, types.ListDataOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionDatumList)

	prb := r.ProboService(ctx, input.OrganizationID)

	pageOrderBy := page.OrderBy[coredata.DatumOrderField]{
		Field:     coredata.DatumOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.DatumOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	noSnapshot := (*gid.GID)(nil)
	datumFilter := coredata.NewDatumFilter(&noSnapshot)
	if input.Filter != nil {
		datumFilter = coredata.NewDatumFilter(&input.Filter.SnapshotID)
	}

	page, err := prb.Data.ListForOrganizationID(ctx, input.OrganizationID, cursor, datumFilter)
	if err != nil {
		panic(fmt.Errorf("cannot list organization data: %w", err))
	}

	return nil, types.NewListDataOutput(page), nil
}

func (r *Resolver) GetDatumTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetDatumInput) (*mcp.CallToolResult, types.GetDatumOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionDatumGet)

	prb := r.ProboService(ctx, input.ID)

	datum, err := prb.Data.Get(ctx, input.ID)
	if err != nil {
		return nil, types.GetDatumOutput{}, fmt.Errorf("failed to get datum: %w", err)
	}

	return nil, types.GetDatumOutput{
		Datum: types.NewDatum(datum),
	}, nil
}

func (r *Resolver) AddDatumTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddDatumInput) (*mcp.CallToolResult, types.AddDatumOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionDatumCreate)

	svc := r.ProboService(ctx, input.OrganizationID)

	datum, err := svc.Data.Create(
		ctx,
		probo.CreateDatumRequest{
			OrganizationID:     input.OrganizationID,
			Name:               input.Name,
			DataClassification: input.DataClassification,
			OwnerID:            input.OwnerID,
			VendorIDs:          input.VendorIds,
		},
	)
	if err != nil {
		return nil, types.AddDatumOutput{}, fmt.Errorf("failed to create datum: %w", err)
	}

	return nil, types.AddDatumOutput{
		Datum: types.NewDatum(datum),
	}, nil
}

func (r *Resolver) UpdateDatumTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateDatumInput) (*mcp.CallToolResult, types.UpdateDatumOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionDatumUpdate)

	svc := r.ProboService(ctx, input.ID)

	datum, err := svc.Data.Update(
		ctx,
		probo.UpdateDatumRequest{
			ID:                 input.ID,
			Name:               input.Name,
			DataClassification: input.DataClassification,
			OwnerID:            input.OwnerID,
			VendorIDs:          input.VendorIds,
		},
	)
	if err != nil {
		return nil, types.UpdateDatumOutput{}, fmt.Errorf("failed to update datum: %w", err)
	}

	return nil, types.UpdateDatumOutput{
		Datum: types.NewDatum(datum),
	}, nil
}

func (r *Resolver) ListNonconformitiesTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListNonconformitiesInput) (*mcp.CallToolResult, types.ListNonconformitiesOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionNonconformityList)

	prb := r.ProboService(ctx, input.OrganizationID)

	pageOrderBy := page.OrderBy[coredata.NonconformityOrderField]{
		Field:     coredata.NonconformityOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.NonconformityOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	noSnapshot := (*gid.GID)(nil)
	nonconformityFilter := coredata.NewNonconformityFilter(&noSnapshot)
	if input.Filter != nil {
		nonconformityFilter = coredata.NewNonconformityFilter(&input.Filter.SnapshotID)
	}

	page, err := prb.Nonconformities.ListForOrganizationID(ctx, input.OrganizationID, cursor, nonconformityFilter)
	if err != nil {
		panic(fmt.Errorf("cannot list organization nonconformities: %w", err))
	}

	return nil, types.NewListNonconformitiesOutput(page), nil
}

func (r *Resolver) GetNonconformityTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetNonconformityInput) (*mcp.CallToolResult, types.GetNonconformityOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionNonconformityGet)

	prb := r.ProboService(ctx, input.ID)

	nonconformity, err := prb.Nonconformities.Get(ctx, input.ID)
	if err != nil {
		return nil, types.GetNonconformityOutput{}, fmt.Errorf("failed to get nonconformity: %w", err)
	}

	return nil, types.GetNonconformityOutput{
		Nonconformity: types.NewNonconformity(nonconformity),
	}, nil
}

func (r *Resolver) AddNonconformityTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddNonconformityInput) (*mcp.CallToolResult, types.AddNonconformityOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionNonconformityCreate)

	svc := r.ProboService(ctx, input.OrganizationID)

	nonconformity, err := svc.Nonconformities.Create(
		ctx,
		&probo.CreateNonconformityRequest{
			OrganizationID:     input.OrganizationID,
			ReferenceID:        input.ReferenceID,
			Description:        input.Description,
			AuditID:            input.AuditID,
			DateIdentified:     input.DateIdentified,
			RootCause:          input.RootCause,
			CorrectiveAction:   input.CorrectiveAction,
			OwnerID:            input.OwnerID,
			DueDate:            input.DueDate,
			Status:             input.Status,
			EffectivenessCheck: input.EffectivenessCheck,
		},
	)
	if err != nil {
		return nil, types.AddNonconformityOutput{}, fmt.Errorf("failed to create nonconformity: %w", err)
	}

	return nil, types.AddNonconformityOutput{
		Nonconformity: types.NewNonconformity(nonconformity),
	}, nil
}

func (r *Resolver) UpdateNonconformityTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateNonconformityInput) (*mcp.CallToolResult, types.UpdateNonconformityOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionNonconformityUpdate)

	svc := r.ProboService(ctx, input.ID)

	nonconformity, err := svc.Nonconformities.Update(
		ctx,
		&probo.UpdateNonconformityRequest{
			ID:                 input.ID,
			ReferenceID:        input.ReferenceID,
			Description:        UnwrapOmittable(input.Description),
			DateIdentified:     UnwrapOmittable(input.DateIdentified),
			RootCause:          input.RootCause,
			CorrectiveAction:   UnwrapOmittable(input.CorrectiveAction),
			OwnerID:            input.OwnerID,
			AuditID:            UnwrapOmittable(input.AuditID),
			DueDate:            UnwrapOmittable(input.DueDate),
			Status:             input.Status,
			EffectivenessCheck: UnwrapOmittable(input.EffectivenessCheck),
		},
	)
	if err != nil {
		return nil, types.UpdateNonconformityOutput{}, fmt.Errorf("failed to update nonconformity: %w", err)
	}

	return nil, types.UpdateNonconformityOutput{
		Nonconformity: types.NewNonconformity(nonconformity),
	}, nil
}

func (r *Resolver) ListObligationsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListObligationsInput) (*mcp.CallToolResult, types.ListObligationsOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionObligationList)

	prb := r.ProboService(ctx, input.OrganizationID)

	pageOrderBy := page.OrderBy[coredata.ObligationOrderField]{
		Field:     coredata.ObligationOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.ObligationOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	noSnapshot := (*gid.GID)(nil)
	obligationFilter := coredata.NewObligationFilter(&noSnapshot)
	if input.Filter != nil {
		obligationFilter = coredata.NewObligationFilter(&input.Filter.SnapshotID)
	}

	page, err := prb.Obligations.ListForOrganizationID(ctx, input.OrganizationID, cursor, obligationFilter)
	if err != nil {
		panic(fmt.Errorf("cannot list organization obligations: %w", err))
	}

	return nil, types.NewListObligationsOutput(page), nil
}

func (r *Resolver) GetObligationTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetObligationInput) (*mcp.CallToolResult, types.GetObligationOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionObligationGet)

	prb := r.ProboService(ctx, input.ID)

	obligation, err := prb.Obligations.Get(ctx, input.ID)
	if err != nil {
		return nil, types.GetObligationOutput{}, fmt.Errorf("failed to get obligation: %w", err)
	}

	return nil, types.GetObligationOutput{
		Obligation: types.NewObligation(obligation),
	}, nil
}

func (r *Resolver) AddObligationTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddObligationInput) (*mcp.CallToolResult, types.AddObligationOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionObligationCreate)

	svc := r.ProboService(ctx, input.OrganizationID)

	obligation, err := svc.Obligations.Create(
		ctx,
		&probo.CreateObligationRequest{
			OrganizationID:         input.OrganizationID,
			Area:                   input.Area,
			Source:                 input.Source,
			Requirement:            input.Requirement,
			ActionsToBeImplemented: input.ActionsToBeImplemented,
			Regulator:              input.Regulator,
			OwnerID:                input.OwnerID,
			LastReviewDate:         input.LastReviewDate,
			DueDate:                input.DueDate,
			Status:                 *input.Status,
			Type:                   *input.Type,
		},
	)
	if err != nil {
		return nil, types.AddObligationOutput{}, fmt.Errorf("failed to create obligation: %w", err)
	}

	return nil, types.AddObligationOutput{
		Obligation: types.NewObligation(obligation),
	}, nil
}

func (r *Resolver) UpdateObligationTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateObligationInput) (*mcp.CallToolResult, types.UpdateObligationOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionObligationUpdate)

	svc := r.ProboService(ctx, input.ID)

	obligation, err := svc.Obligations.Update(
		ctx,
		&probo.UpdateObligationRequest{
			ID:                     input.ID,
			Area:                   UnwrapOmittable(input.Area),
			Source:                 UnwrapOmittable(input.Source),
			Requirement:            UnwrapOmittable(input.Requirement),
			ActionsToBeImplemented: UnwrapOmittable(input.ActionsToBeImplemented),
			Regulator:              UnwrapOmittable(input.Regulator),
			OwnerID:                input.OwnerID,
			LastReviewDate:         UnwrapOmittable(input.LastReviewDate),
			DueDate:                UnwrapOmittable(input.DueDate),
			Status:                 input.Status,
			Type:                   input.Type,
		},
	)
	if err != nil {
		return nil, types.UpdateObligationOutput{}, fmt.Errorf("failed to update obligation: %w", err)
	}

	return nil, types.UpdateObligationOutput{
		Obligation: types.NewObligation(obligation),
	}, nil
}

func (r *Resolver) ListProcessingActivitiesTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListProcessingActivitiesInput) (*mcp.CallToolResult, types.ListProcessingActivitiesOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionProcessingActivityList)

	prb := r.ProboService(ctx, input.OrganizationID)

	pageOrderBy := page.OrderBy[coredata.ProcessingActivityOrderField]{
		Field:     coredata.ProcessingActivityOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.ProcessingActivityOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	noSnapshot := (*gid.GID)(nil)
	filter := coredata.NewProcessingActivityFilter(&noSnapshot)
	if input.Filter != nil {
		filter = coredata.NewProcessingActivityFilter(&input.Filter.SnapshotID)
	}

	page, err := prb.ProcessingActivities.ListForOrganizationID(ctx, input.OrganizationID, cursor, filter)
	if err != nil {
		panic(fmt.Errorf("cannot list organization processing activities: %w", err))
	}

	return nil, types.NewListProcessingActivitiesOutput(page), nil
}

func (r *Resolver) GetProcessingActivityTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetProcessingActivityInput) (*mcp.CallToolResult, types.GetProcessingActivityOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionProcessingActivityGet)

	prb := r.ProboService(ctx, input.ID)

	processingActivity, err := prb.ProcessingActivities.Get(ctx, input.ID)
	if err != nil {
		return nil, types.GetProcessingActivityOutput{}, fmt.Errorf("failed to get processing activity: %w", err)
	}

	return nil, types.GetProcessingActivityOutput{
		ProcessingActivity: types.NewProcessingActivity(processingActivity),
	}, nil
}

func (r *Resolver) AddProcessingActivityTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddProcessingActivityInput) (*mcp.CallToolResult, types.AddProcessingActivityOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionProcessingActivityCreate)

	svc := r.ProboService(ctx, input.OrganizationID)

	processingActivity, err := svc.ProcessingActivities.Create(
		ctx,
		&probo.CreateProcessingActivityRequest{
			OrganizationID:                       input.OrganizationID,
			Name:                                 input.Name,
			Purpose:                              input.Purpose,
			DataSubjectCategory:                  input.DataSubjectCategory,
			PersonalDataCategory:                 input.PersonalDataCategory,
			SpecialOrCriminalData:                input.SpecialOrCriminalData,
			ConsentEvidenceLink:                  input.ConsentEvidenceLink,
			LawfulBasis:                          input.LawfulBasis,
			Recipients:                           input.Recipients,
			Location:                             input.Location,
			InternationalTransfers:               input.InternationalTransfers,
			TransferSafeguard:                    input.TransferSafeguard,
			RetentionPeriod:                      input.RetentionPeriod,
			SecurityMeasures:                     input.SecurityMeasures,
			DataProtectionImpactAssessmentNeeded: input.DataProtectionImpactAssessmentNeeded,
			TransferImpactAssessmentNeeded:       input.TransferImpactAssessmentNeeded,
			LastReviewDate:                       input.LastReviewDate,
			NextReviewDate:                       input.NextReviewDate,
			Role:                                 input.Role,
			DataProtectionOfficerID:              input.DataProtectionOfficerID,
			VendorIDs:                            input.VendorIds,
		},
	)
	if err != nil {
		return nil, types.AddProcessingActivityOutput{}, fmt.Errorf("failed to create processing activity: %w", err)
	}

	return nil, types.AddProcessingActivityOutput{
		ProcessingActivity: types.NewProcessingActivity(processingActivity),
	}, nil
}

func (r *Resolver) UpdateProcessingActivityTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateProcessingActivityInput) (*mcp.CallToolResult, types.UpdateProcessingActivityOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionProcessingActivityUpdate)

	svc := r.ProboService(ctx, input.ID)

	var vendorIDs *[]gid.GID
	if input.VendorIds != nil {
		vendorIDs = &input.VendorIds
	}

	processingActivity, err := svc.ProcessingActivities.Update(
		ctx,
		&probo.UpdateProcessingActivityRequest{
			ID:                                   input.ID,
			Name:                                 input.Name,
			Purpose:                              UnwrapOmittable(input.Purpose),
			DataSubjectCategory:                  UnwrapOmittable(input.DataSubjectCategory),
			PersonalDataCategory:                 UnwrapOmittable(input.PersonalDataCategory),
			SpecialOrCriminalData:                input.SpecialOrCriminalData,
			ConsentEvidenceLink:                  UnwrapOmittable(input.ConsentEvidenceLink),
			LawfulBasis:                          input.LawfulBasis,
			Recipients:                           UnwrapOmittable(input.Recipients),
			Location:                             UnwrapOmittable(input.Location),
			InternationalTransfers:               input.InternationalTransfers,
			TransferSafeguard:                    UnwrapOmittable(input.TransferSafeguard),
			RetentionPeriod:                      UnwrapOmittable(input.RetentionPeriod),
			SecurityMeasures:                     UnwrapOmittable(input.SecurityMeasures),
			DataProtectionImpactAssessmentNeeded: input.DataProtectionImpactAssessmentNeeded,
			TransferImpactAssessmentNeeded:       input.TransferImpactAssessmentNeeded,
			LastReviewDate:                       UnwrapOmittable(input.LastReviewDate),
			NextReviewDate:                       UnwrapOmittable(input.NextReviewDate),
			Role:                                 input.Role,
			DataProtectionOfficerID:              UnwrapOmittable(input.DataProtectionOfficerID),
			VendorIDs:                            vendorIDs,
		},
	)
	if err != nil {
		return nil, types.UpdateProcessingActivityOutput{}, fmt.Errorf("failed to update processing activity: %w", err)
	}

	return nil, types.UpdateProcessingActivityOutput{
		ProcessingActivity: types.NewProcessingActivity(processingActivity),
	}, nil
}

func (r *Resolver) DeleteProcessingActivityTool(ctx context.Context, req *mcp.CallToolRequest, input *types.DeleteProcessingActivityInput) (*mcp.CallToolResult, types.DeleteProcessingActivityOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionProcessingActivityDelete)

	svc := r.ProboService(ctx, input.ID)

	err := svc.ProcessingActivities.Delete(ctx, input.ID)
	if err != nil {
		return nil, types.DeleteProcessingActivityOutput{}, fmt.Errorf("failed to delete processing activity: %w", err)
	}

	return nil, types.DeleteProcessingActivityOutput{
		DeletedProcessingActivityID: input.ID,
	}, nil
}

func (r *Resolver) ListDataProtectionImpactAssessmentsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListDataProtectionImpactAssessmentsInput) (*mcp.CallToolResult, types.ListDataProtectionImpactAssessmentsOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionDataProtectionImpactAssessmentList)

	prb := r.ProboService(ctx, input.OrganizationID)

	pageOrderBy := page.OrderBy[coredata.DataProtectionImpactAssessmentOrderField]{
		Field:     coredata.DataProtectionImpactAssessmentOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.DataProtectionImpactAssessmentOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	noSnapshot := (*gid.GID)(nil)
	filter := coredata.NewDataProtectionImpactAssessmentFilter(&noSnapshot)
	if input.Filter != nil {
		filter = coredata.NewDataProtectionImpactAssessmentFilter(&input.Filter.SnapshotID)
	}

	page, err := prb.DataProtectionImpactAssessments.ListForOrganizationID(ctx, input.OrganizationID, cursor, filter)
	if err != nil {
		panic(fmt.Errorf("cannot list organization data protection impact assessments: %w", err))
	}

	return nil, types.NewListDataProtectionImpactAssessmentsOutput(page), nil
}

func (r *Resolver) GetDataProtectionImpactAssessmentTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetDataProtectionImpactAssessmentInput) (*mcp.CallToolResult, types.GetDataProtectionImpactAssessmentOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionDataProtectionImpactAssessmentGet)

	prb := r.ProboService(ctx, input.ID)

	dpia, err := prb.DataProtectionImpactAssessments.Get(ctx, input.ID)
	if err != nil {
		return nil, types.GetDataProtectionImpactAssessmentOutput{}, fmt.Errorf("failed to get data protection impact assessment: %w", err)
	}

	return nil, types.GetDataProtectionImpactAssessmentOutput{
		DataProtectionImpactAssessment: types.NewDataProtectionImpactAssessment(dpia),
	}, nil
}

func (r *Resolver) AddDataProtectionImpactAssessmentTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddDataProtectionImpactAssessmentInput) (*mcp.CallToolResult, types.AddDataProtectionImpactAssessmentOutput, error) {
	r.MustAuthorize(ctx, input.ProcessingActivityID, probo.ActionDataProtectionImpactAssessmentCreate)

	svc := r.ProboService(ctx, input.ProcessingActivityID)

	dpia, err := svc.DataProtectionImpactAssessments.Create(
		ctx,
		&probo.CreateDataProtectionImpactAssessmentRequest{
			ProcessingActivityID:        input.ProcessingActivityID,
			Description:                 input.Description,
			NecessityAndProportionality: input.NecessityAndProportionality,
			PotentialRisk:               input.PotentialRisk,
			Mitigations:                 input.Mitigations,
			ResidualRisk:                input.ResidualRisk,
		},
	)
	if err != nil {
		return nil, types.AddDataProtectionImpactAssessmentOutput{}, fmt.Errorf("failed to create data protection impact assessment: %w", err)
	}

	return nil, types.AddDataProtectionImpactAssessmentOutput{
		DataProtectionImpactAssessment: types.NewDataProtectionImpactAssessment(dpia),
	}, nil
}

func (r *Resolver) UpdateDataProtectionImpactAssessmentTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateDataProtectionImpactAssessmentInput) (*mcp.CallToolResult, types.UpdateDataProtectionImpactAssessmentOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionDataProtectionImpactAssessmentUpdate)

	svc := r.ProboService(ctx, input.ID)

	dpia, err := svc.DataProtectionImpactAssessments.Update(
		ctx,
		&probo.UpdateDataProtectionImpactAssessmentRequest{
			ID:                          input.ID,
			Description:                 UnwrapOmittable(input.Description),
			NecessityAndProportionality: UnwrapOmittable(input.NecessityAndProportionality),
			PotentialRisk:               UnwrapOmittable(input.PotentialRisk),
			Mitigations:                 UnwrapOmittable(input.Mitigations),
			ResidualRisk:                input.ResidualRisk,
		},
	)
	if err != nil {
		return nil, types.UpdateDataProtectionImpactAssessmentOutput{}, fmt.Errorf("failed to update data protection impact assessment: %w", err)
	}

	return nil, types.UpdateDataProtectionImpactAssessmentOutput{
		DataProtectionImpactAssessment: types.NewDataProtectionImpactAssessment(dpia),
	}, nil
}

func (r *Resolver) ListTransferImpactAssessmentsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListTransferImpactAssessmentsInput) (*mcp.CallToolResult, types.ListTransferImpactAssessmentsOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionTransferImpactAssessmentList)

	prb := r.ProboService(ctx, input.OrganizationID)

	pageOrderBy := page.OrderBy[coredata.TransferImpactAssessmentOrderField]{
		Field:     coredata.TransferImpactAssessmentOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.TransferImpactAssessmentOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	noSnapshot := (*gid.GID)(nil)
	filter := coredata.NewTransferImpactAssessmentFilter(&noSnapshot)
	if input.Filter != nil {
		filter = coredata.NewTransferImpactAssessmentFilter(&input.Filter.SnapshotID)
	}

	page, err := prb.TransferImpactAssessments.ListForOrganizationID(ctx, input.OrganizationID, cursor, filter)
	if err != nil {
		panic(fmt.Errorf("cannot list organization transfer impact assessments: %w", err))
	}

	return nil, types.NewListTransferImpactAssessmentsOutput(page), nil
}

func (r *Resolver) GetTransferImpactAssessmentTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetTransferImpactAssessmentInput) (*mcp.CallToolResult, types.GetTransferImpactAssessmentOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionTransferImpactAssessmentGet)

	prb := r.ProboService(ctx, input.ID)

	tia, err := prb.TransferImpactAssessments.Get(ctx, input.ID)
	if err != nil {
		return nil, types.GetTransferImpactAssessmentOutput{}, fmt.Errorf("failed to get transfer impact assessment: %w", err)
	}

	return nil, types.GetTransferImpactAssessmentOutput{
		TransferImpactAssessment: types.NewTransferImpactAssessment(tia),
	}, nil
}

func (r *Resolver) AddTransferImpactAssessmentTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddTransferImpactAssessmentInput) (*mcp.CallToolResult, types.AddTransferImpactAssessmentOutput, error) {
	r.MustAuthorize(ctx, input.ProcessingActivityID, probo.ActionTransferImpactAssessmentCreate)

	svc := r.ProboService(ctx, input.ProcessingActivityID)

	tia, err := svc.TransferImpactAssessments.Create(
		ctx,
		&probo.CreateTransferImpactAssessmentRequest{
			ProcessingActivityID:  input.ProcessingActivityID,
			DataSubjects:          input.DataSubjects,
			LegalMechanism:        input.LegalMechanism,
			Transfer:              input.Transfer,
			LocalLawRisk:          input.LocalLawRisk,
			SupplementaryMeasures: input.SupplementaryMeasures,
		},
	)
	if err != nil {
		return nil, types.AddTransferImpactAssessmentOutput{}, fmt.Errorf("failed to create transfer impact assessment: %w", err)
	}

	return nil, types.AddTransferImpactAssessmentOutput{
		TransferImpactAssessment: types.NewTransferImpactAssessment(tia),
	}, nil
}

func (r *Resolver) UpdateTransferImpactAssessmentTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateTransferImpactAssessmentInput) (*mcp.CallToolResult, types.UpdateTransferImpactAssessmentOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionTransferImpactAssessmentUpdate)

	svc := r.ProboService(ctx, input.ID)

	tia, err := svc.TransferImpactAssessments.Update(
		ctx,
		&probo.UpdateTransferImpactAssessmentRequest{
			ID:                    input.ID,
			DataSubjects:          UnwrapOmittable(input.DataSubjects),
			LegalMechanism:        UnwrapOmittable(input.LegalMechanism),
			Transfer:              UnwrapOmittable(input.Transfer),
			LocalLawRisk:          UnwrapOmittable(input.LocalLawRisk),
			SupplementaryMeasures: UnwrapOmittable(input.SupplementaryMeasures),
		},
	)
	if err != nil {
		return nil, types.UpdateTransferImpactAssessmentOutput{}, fmt.Errorf("failed to update transfer impact assessment: %w", err)
	}

	return nil, types.UpdateTransferImpactAssessmentOutput{
		TransferImpactAssessment: types.NewTransferImpactAssessment(tia),
	}, nil
}

func (r *Resolver) DeleteTransferImpactAssessmentTool(ctx context.Context, req *mcp.CallToolRequest, input *types.DeleteTransferImpactAssessmentInput) (*mcp.CallToolResult, types.DeleteTransferImpactAssessmentOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionTransferImpactAssessmentDelete)

	svc := r.ProboService(ctx, input.ID)

	err := svc.TransferImpactAssessments.Delete(ctx, input.ID)
	if err != nil {
		return nil, types.DeleteTransferImpactAssessmentOutput{}, fmt.Errorf("failed to delete transfer impact assessment: %w", err)
	}

	return nil, types.DeleteTransferImpactAssessmentOutput{
		DeletedTransferImpactAssessmentID: input.ID,
	}, nil
}

func (r *Resolver) ListContinualImprovementsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListContinualImprovementsInput) (*mcp.CallToolResult, types.ListContinualImprovementsOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionContinualImprovementList)

	prb := r.ProboService(ctx, input.OrganizationID)

	pageOrderBy := page.OrderBy[coredata.ContinualImprovementOrderField]{
		Field:     coredata.ContinualImprovementOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.ContinualImprovementOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	noSnapshot := (*gid.GID)(nil)
	continualImprovementFilter := coredata.NewContinualImprovementFilter(&noSnapshot)
	if input.Filter != nil {
		continualImprovementFilter = coredata.NewContinualImprovementFilter(&input.Filter.SnapshotID)
	}

	page, err := prb.ContinualImprovements.ListForOrganizationID(ctx, input.OrganizationID, cursor, continualImprovementFilter)
	if err != nil {
		panic(fmt.Errorf("cannot list organization continual improvements: %w", err))
	}

	return nil, types.NewListContinualImprovementsOutput(page), nil
}

func (r *Resolver) GetContinualImprovementTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetContinualImprovementInput) (*mcp.CallToolResult, types.GetContinualImprovementOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionContinualImprovementGet)

	prb := r.ProboService(ctx, input.ID)

	continualImprovement, err := prb.ContinualImprovements.Get(ctx, input.ID)
	if err != nil {
		return nil, types.GetContinualImprovementOutput{}, fmt.Errorf("failed to get continual improvement: %w", err)
	}

	return nil, types.GetContinualImprovementOutput{
		ContinualImprovement: types.NewContinualImprovement(continualImprovement),
	}, nil
}

func (r *Resolver) AddContinualImprovementTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddContinualImprovementInput) (*mcp.CallToolResult, types.AddContinualImprovementOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionContinualImprovementCreate)

	svc := r.ProboService(ctx, input.OrganizationID)

	continualImprovement, err := svc.ContinualImprovements.Create(
		ctx,
		&probo.CreateContinualImprovementRequest{
			OrganizationID: input.OrganizationID,
			ReferenceID:    input.ReferenceID,
			Description:    input.Description,
			Source:         input.Source,
			OwnerID:        input.OwnerID,
			TargetDate:     input.TargetDate,
			Status:         input.Status,
			Priority:       input.Priority,
		},
	)
	if err != nil {
		return nil, types.AddContinualImprovementOutput{}, fmt.Errorf("failed to create continual improvement: %w", err)
	}

	return nil, types.AddContinualImprovementOutput{
		ContinualImprovement: types.NewContinualImprovement(continualImprovement),
	}, nil
}

func (r *Resolver) UpdateContinualImprovementTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateContinualImprovementInput) (*mcp.CallToolResult, types.UpdateContinualImprovementOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionContinualImprovementUpdate)

	svc := r.ProboService(ctx, input.ID)

	continualImprovement, err := svc.ContinualImprovements.Update(
		ctx,
		&probo.UpdateContinualImprovementRequest{
			ID:          input.ID,
			ReferenceID: input.ReferenceID,
			Description: UnwrapOmittable(input.Description),
			Source:      UnwrapOmittable(input.Source),
			OwnerID:     input.OwnerID,
			TargetDate:  UnwrapOmittable(input.TargetDate),
			Status:      input.Status,
			Priority:    input.Priority,
		},
	)
	if err != nil {
		return nil, types.UpdateContinualImprovementOutput{}, fmt.Errorf("failed to update continual improvement: %w", err)
	}

	return nil, types.UpdateContinualImprovementOutput{
		ContinualImprovement: types.NewContinualImprovement(continualImprovement),
	}, nil
}

func (r *Resolver) ListAuditsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListAuditsInput) (*mcp.CallToolResult, types.ListAuditsOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionAuditList)

	prb := r.ProboService(ctx, input.OrganizationID)

	pageOrderBy := page.OrderBy[coredata.AuditOrderField]{
		Field:     coredata.AuditOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.AuditOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	page, err := prb.Audits.ListForOrganizationID(ctx, input.OrganizationID, cursor)
	if err != nil {
		panic(fmt.Errorf("cannot list organization audits: %w", err))
	}

	return nil, types.NewListAuditsOutput(page), nil
}

func (r *Resolver) GetAuditTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetAuditInput) (*mcp.CallToolResult, types.GetAuditOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionAuditGet)

	prb := r.ProboService(ctx, input.ID)

	audit, err := prb.Audits.Get(ctx, input.ID)
	if err != nil {
		return nil, types.GetAuditOutput{}, fmt.Errorf("failed to get audit: %w", err)
	}

	return nil, types.GetAuditOutput{
		Audit: types.NewAudit(audit),
	}, nil
}

func (r *Resolver) AddAuditTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddAuditInput) (*mcp.CallToolResult, types.AddAuditOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionAuditCreate)

	svc := r.ProboService(ctx, input.OrganizationID)

	audit, err := svc.Audits.Create(
		ctx,
		&probo.CreateAuditRequest{
			OrganizationID: input.OrganizationID,
			Name:           input.Name,
			ValidFrom:      input.ValidFrom,
			ValidUntil:     input.ValidUntil,
			State:          input.State,
			FrameworkID:    input.FrameworkID,
		},
	)
	if err != nil {
		return nil, types.AddAuditOutput{}, fmt.Errorf("failed to create audit: %w", err)
	}

	return nil, types.AddAuditOutput{
		Audit: types.NewAudit(audit),
	}, nil
}

func (r *Resolver) UpdateAuditTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateAuditInput) (*mcp.CallToolResult, types.UpdateAuditOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionAuditUpdate)

	svc := r.ProboService(ctx, input.ID)

	audit, err := svc.Audits.Update(
		ctx,
		&probo.UpdateAuditRequest{
			ID:                    input.ID,
			Name:                  UnwrapOmittable(input.Name),
			ValidFrom:             input.ValidFrom,
			ValidUntil:            input.ValidUntil,
			State:                 input.State,
			TrustCenterVisibility: input.TrustCenterVisibility,
		},
	)
	if err != nil {
		return nil, types.UpdateAuditOutput{}, fmt.Errorf("failed to update audit: %w", err)
	}

	return nil, types.UpdateAuditOutput{
		Audit: types.NewAudit(audit),
	}, nil
}

func (r *Resolver) ListControlsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListControlsInput) (*mcp.CallToolResult, types.ListControlsOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionControlList)

	prb := r.ProboService(ctx, input.OrganizationID)

	pageOrderBy := page.OrderBy[coredata.ControlOrderField]{
		Field:     coredata.ControlOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.ControlOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	var controlFilter = coredata.NewControlFilter(nil)
	if input.Filter != nil {
		controlFilter = coredata.NewControlFilter(input.Filter.Query)
	}

	page, err := prb.Controls.ListForOrganizationID(ctx, input.OrganizationID, cursor, controlFilter)
	if err != nil {
		panic(fmt.Errorf("cannot list organization controls: %w", err))
	}

	return nil, types.NewListControlsOutput(page), nil
}

func (r *Resolver) GetControlTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetControlInput) (*mcp.CallToolResult, types.GetControlOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionControlGet)

	prb := r.ProboService(ctx, input.ID)

	control, err := prb.Controls.Get(ctx, input.ID)
	if err != nil {
		return nil, types.GetControlOutput{}, fmt.Errorf("failed to get control: %w", err)
	}

	return nil, types.GetControlOutput{
		Control: types.NewControl(control),
	}, nil
}

func (r *Resolver) AddControlTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddControlInput) (*mcp.CallToolResult, types.AddControlOutput, error) {
	r.MustAuthorize(ctx, input.FrameworkID, probo.ActionControlCreate)

	svc := r.ProboService(ctx, input.FrameworkID)

	control, err := svc.Controls.Create(
		ctx,
		probo.CreateControlRequest{
			FrameworkID:  input.FrameworkID,
			Name:         input.Name,
			Description:  input.Description,
			SectionTitle: input.SectionTitle,
		},
	)
	if err != nil {
		return nil, types.AddControlOutput{}, fmt.Errorf("failed to create control: %w", err)
	}

	return nil, types.AddControlOutput{
		Control: types.NewControl(control),
	}, nil
}

func (r *Resolver) UpdateControlTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateControlInput) (*mcp.CallToolResult, types.UpdateControlOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionControlUpdate)

	svc := r.ProboService(ctx, input.ID)

	control, err := svc.Controls.Update(
		ctx,
		probo.UpdateControlRequest{
			ID:           input.ID,
			Name:         input.Name,
			Description:  UnwrapOmittable(input.Description),
			SectionTitle: input.SectionTitle,
		},
	)
	if err != nil {
		return nil, types.UpdateControlOutput{}, fmt.Errorf("failed to update control: %w", err)
	}

	return nil, types.UpdateControlOutput{
		Control: types.NewControl(control),
	}, nil
}

func (r *Resolver) LinkControlTool(ctx context.Context, req *mcp.CallToolRequest, input *types.LinkControlInput) (*mcp.CallToolResult, types.LinkControlOutput, error) {
	svc := r.ProboService(ctx, input.ControlID)

	switch input.ResourceID.EntityType() {
	case coredata.MeasureEntityType:
		r.MustAuthorize(ctx, input.ControlID, probo.ActionControlMeasureMappingCreate)
		if _, _, err := svc.Controls.CreateMeasureMapping(ctx, input.ControlID, input.ResourceID); err != nil {
			return nil, types.LinkControlOutput{}, fmt.Errorf("failed to link control to measure: %w", err)
		}
	case coredata.DocumentEntityType:
		r.MustAuthorize(ctx, input.ControlID, probo.ActionControlDocumentMappingCreate)
		if _, _, err := svc.Controls.CreateDocumentMapping(ctx, input.ControlID, input.ResourceID); err != nil {
			return nil, types.LinkControlOutput{}, fmt.Errorf("failed to link control to document: %w", err)
		}
	case coredata.AuditEntityType:
		r.MustAuthorize(ctx, input.ControlID, probo.ActionControlAuditMappingCreate)
		if _, _, err := svc.Controls.CreateAuditMapping(ctx, input.ControlID, input.ResourceID); err != nil {
			return nil, types.LinkControlOutput{}, fmt.Errorf("failed to link control to audit: %w", err)
		}
	case coredata.SnapshotEntityType:
		r.MustAuthorize(ctx, input.ControlID, probo.ActionControlSnapshotMappingCreate)
		if _, _, err := svc.Controls.CreateSnapshotMapping(ctx, input.ControlID, input.ResourceID); err != nil {
			return nil, types.LinkControlOutput{}, fmt.Errorf("failed to link control to snapshot: %w", err)
		}
	case coredata.ObligationEntityType:
		r.MustAuthorize(ctx, input.ControlID, probo.ActionControlObligationMappingCreate)
		if _, _, err := svc.Controls.CreateObligationMapping(ctx, input.ControlID, input.ResourceID); err != nil {
			return nil, types.LinkControlOutput{}, fmt.Errorf("failed to link control to obligation: %w", err)
		}
	default:
		return nil, types.LinkControlOutput{}, fmt.Errorf("unsupported resource type for control linking: entity type %d", input.ResourceID.EntityType())
	}

	return nil, types.LinkControlOutput{}, nil
}

func (r *Resolver) UnlinkControlTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UnlinkControlInput) (*mcp.CallToolResult, types.UnlinkControlOutput, error) {
	svc := r.ProboService(ctx, input.ControlID)

	switch input.ResourceID.EntityType() {
	case coredata.MeasureEntityType:
		r.MustAuthorize(ctx, input.ControlID, probo.ActionControlMeasureMappingDelete)
		if _, _, err := svc.Controls.DeleteMeasureMapping(ctx, input.ControlID, input.ResourceID); err != nil {
			return nil, types.UnlinkControlOutput{}, fmt.Errorf("failed to unlink control from measure: %w", err)
		}
	case coredata.DocumentEntityType:
		r.MustAuthorize(ctx, input.ControlID, probo.ActionControlDocumentMappingDelete)
		if _, _, err := svc.Controls.DeleteDocumentMapping(ctx, input.ControlID, input.ResourceID); err != nil {
			return nil, types.UnlinkControlOutput{}, fmt.Errorf("failed to unlink control from document: %w", err)
		}
	case coredata.AuditEntityType:
		r.MustAuthorize(ctx, input.ControlID, probo.ActionControlAuditMappingDelete)
		if _, _, err := svc.Controls.DeleteAuditMapping(ctx, input.ControlID, input.ResourceID); err != nil {
			return nil, types.UnlinkControlOutput{}, fmt.Errorf("failed to unlink control from audit: %w", err)
		}
	case coredata.SnapshotEntityType:
		r.MustAuthorize(ctx, input.ControlID, probo.ActionControlSnapshotMappingDelete)
		if _, _, err := svc.Controls.DeleteSnapshotMapping(ctx, input.ControlID, input.ResourceID); err != nil {
			return nil, types.UnlinkControlOutput{}, fmt.Errorf("failed to unlink control from snapshot: %w", err)
		}
	case coredata.ObligationEntityType:
		r.MustAuthorize(ctx, input.ControlID, probo.ActionControlObligationMappingDelete)
		if _, _, err := svc.Controls.DeleteObligationMapping(ctx, input.ControlID, input.ResourceID); err != nil {
			return nil, types.UnlinkControlOutput{}, fmt.Errorf("failed to unlink control from obligation: %w", err)
		}
	default:
		return nil, types.UnlinkControlOutput{}, fmt.Errorf("unsupported resource type for control unlinking: entity type %d", input.ResourceID.EntityType())
	}

	return nil, types.UnlinkControlOutput{}, nil
}

func (r *Resolver) ListControlObligationsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListControlObligationsInput) (*mcp.CallToolResult, types.ListControlObligationsOutput, error) {
	r.MustAuthorize(ctx, input.ControlID, probo.ActionControlGet)

	prb := r.ProboService(ctx, input.ControlID)

	pageOrderBy := page.OrderBy[coredata.ObligationOrderField]{
		Field:     coredata.ObligationOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.ObligationOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	obligationPage, err := prb.Obligations.ListForControlID(ctx, input.ControlID, cursor, coredata.NewObligationFilter(nil))
	if err != nil {
		return nil, types.ListControlObligationsOutput{}, fmt.Errorf("failed to list control obligations: %w", err)
	}

	return nil, types.NewListControlObligationsOutput(obligationPage), nil
}

func (r *Resolver) ListControlMeasuresTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListControlMeasuresInput) (*mcp.CallToolResult, types.ListControlMeasuresOutput, error) {
	r.MustAuthorize(ctx, input.ControlID, probo.ActionControlGet)

	prb := r.ProboService(ctx, input.ControlID)

	pageOrderBy := page.OrderBy[coredata.MeasureOrderField]{
		Field:     coredata.MeasureOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.MeasureOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	measurePage, err := prb.Measures.ListForControlID(ctx, input.ControlID, cursor, coredata.NewMeasureFilter(nil, nil))
	if err != nil {
		return nil, types.ListControlMeasuresOutput{}, fmt.Errorf("failed to list control measures: %w", err)
	}

	return nil, types.NewListControlMeasuresOutput(measurePage), nil
}

func (r *Resolver) ListControlDocumentsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListControlDocumentsInput) (*mcp.CallToolResult, types.ListControlDocumentsOutput, error) {
	r.MustAuthorize(ctx, input.ControlID, probo.ActionControlGet)

	prb := r.ProboService(ctx, input.ControlID)

	pageOrderBy := page.OrderBy[coredata.DocumentOrderField]{
		Field:     coredata.DocumentOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.DocumentOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	docPage, err := prb.Documents.ListForControlID(ctx, input.ControlID, cursor, coredata.NewDocumentFilter(nil))
	if err != nil {
		return nil, types.ListControlDocumentsOutput{}, fmt.Errorf("failed to list control documents: %w", err)
	}

	approverIDsMap := make(map[gid.GID][]gid.GID)
	for _, d := range docPage.Data {
		approverPage, err := prb.Documents.ListApprovers(ctx, d.ID, allApproversCursor())
		if err != nil {
			return nil, types.ListControlDocumentsOutput{}, fmt.Errorf("failed to list document approvers: %w", err)
		}
		approverIDsMap[d.ID] = profileIDs(approverPage)
	}

	return nil, types.NewListControlDocumentsOutput(docPage, approverIDsMap), nil
}

func (r *Resolver) ListControlAuditsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListControlAuditsInput) (*mcp.CallToolResult, types.ListControlAuditsOutput, error) {
	r.MustAuthorize(ctx, input.ControlID, probo.ActionControlGet)

	prb := r.ProboService(ctx, input.ControlID)

	pageOrderBy := page.OrderBy[coredata.AuditOrderField]{
		Field:     coredata.AuditOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.AuditOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	auditPage, err := prb.Audits.ListForControlID(ctx, input.ControlID, cursor)
	if err != nil {
		return nil, types.ListControlAuditsOutput{}, fmt.Errorf("failed to list control audits: %w", err)
	}

	return nil, types.NewListControlAuditsOutput(auditPage), nil
}

func (r *Resolver) ListControlSnapshotsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListControlSnapshotsInput) (*mcp.CallToolResult, types.ListControlSnapshotsOutput, error) {
	r.MustAuthorize(ctx, input.ControlID, probo.ActionControlGet)

	prb := r.ProboService(ctx, input.ControlID)

	pageOrderBy := page.OrderBy[coredata.SnapshotOrderField]{
		Field:     coredata.SnapshotOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.SnapshotOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	snapshotPage, err := prb.Snapshots.ListForControlID(ctx, input.ControlID, cursor)
	if err != nil {
		return nil, types.ListControlSnapshotsOutput{}, fmt.Errorf("failed to list control snapshots: %w", err)
	}

	return nil, types.NewListControlSnapshotsOutput(snapshotPage), nil
}

func (r *Resolver) ListRiskObligationsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListRiskObligationsInput) (*mcp.CallToolResult, types.ListRiskObligationsOutput, error) {
	r.MustAuthorize(ctx, input.RiskID, probo.ActionRiskGet)

	prb := r.ProboService(ctx, input.RiskID)

	pageOrderBy := page.OrderBy[coredata.ObligationOrderField]{
		Field:     coredata.ObligationOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.ObligationOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	obligationPage, err := prb.Obligations.ListForRiskID(ctx, input.RiskID, cursor, coredata.NewObligationFilter(nil))
	if err != nil {
		return nil, types.ListRiskObligationsOutput{}, fmt.Errorf("failed to list risk obligations: %w", err)
	}

	return nil, types.NewListRiskObligationsOutput(obligationPage), nil
}

func (r *Resolver) LinkRiskTool(ctx context.Context, req *mcp.CallToolRequest, input *types.LinkRiskInput) (*mcp.CallToolResult, types.LinkRiskOutput, error) {
	svc := r.ProboService(ctx, input.RiskID)

	switch input.ResourceID.EntityType() {
	case coredata.DocumentEntityType:
		r.MustAuthorize(ctx, input.RiskID, probo.ActionRiskDocumentMappingCreate)
		if _, _, err := svc.Risks.CreateDocumentMapping(ctx, input.RiskID, input.ResourceID); err != nil {
			return nil, types.LinkRiskOutput{}, fmt.Errorf("failed to link risk to document: %w", err)
		}
	case coredata.MeasureEntityType:
		r.MustAuthorize(ctx, input.RiskID, probo.ActionRiskMeasureMappingCreate)
		if _, _, err := svc.Risks.CreateMeasureMapping(ctx, input.RiskID, input.ResourceID); err != nil {
			return nil, types.LinkRiskOutput{}, fmt.Errorf("failed to link risk to measure: %w", err)
		}
	case coredata.ObligationEntityType:
		r.MustAuthorize(ctx, input.RiskID, probo.ActionRiskObligationMappingCreate)
		if _, _, err := svc.Risks.CreateObligationMapping(ctx, input.RiskID, input.ResourceID); err != nil {
			return nil, types.LinkRiskOutput{}, fmt.Errorf("failed to link risk to obligation: %w", err)
		}
	default:
		return nil, types.LinkRiskOutput{}, fmt.Errorf("unsupported resource type for risk linking: entity type %d", input.ResourceID.EntityType())
	}

	return nil, types.LinkRiskOutput{}, nil
}

func (r *Resolver) UnlinkRiskTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UnlinkRiskInput) (*mcp.CallToolResult, types.UnlinkRiskOutput, error) {
	svc := r.ProboService(ctx, input.RiskID)

	switch input.ResourceID.EntityType() {
	case coredata.DocumentEntityType:
		r.MustAuthorize(ctx, input.RiskID, probo.ActionRiskDocumentMappingDelete)
		if _, _, err := svc.Risks.DeleteDocumentMapping(ctx, input.RiskID, input.ResourceID); err != nil {
			return nil, types.UnlinkRiskOutput{}, fmt.Errorf("failed to unlink risk from document: %w", err)
		}
	case coredata.MeasureEntityType:
		r.MustAuthorize(ctx, input.RiskID, probo.ActionRiskMeasureMappingDelete)
		if _, _, err := svc.Risks.DeleteMeasureMapping(ctx, input.RiskID, input.ResourceID); err != nil {
			return nil, types.UnlinkRiskOutput{}, fmt.Errorf("failed to unlink risk from measure: %w", err)
		}
	case coredata.ObligationEntityType:
		r.MustAuthorize(ctx, input.RiskID, probo.ActionRiskObligationMappingDelete)
		if _, _, err := svc.Risks.DeleteObligationMapping(ctx, input.RiskID, input.ResourceID); err != nil {
			return nil, types.UnlinkRiskOutput{}, fmt.Errorf("failed to unlink risk from obligation: %w", err)
		}
	default:
		return nil, types.UnlinkRiskOutput{}, fmt.Errorf("unsupported resource type for risk unlinking: entity type %d", input.ResourceID.EntityType())
	}

	return nil, types.UnlinkRiskOutput{}, nil
}

func (r *Resolver) ListTasksTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListTasksInput) (*mcp.CallToolResult, types.ListTasksOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionTaskList)

	prb := r.ProboService(ctx, input.OrganizationID)

	pageOrderBy := page.OrderBy[coredata.TaskOrderField]{
		Field:     coredata.TaskOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.TaskOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	page, err := prb.Tasks.ListForOrganizationID(ctx, input.OrganizationID, cursor)
	if err != nil {
		panic(fmt.Errorf("cannot list organization tasks: %w", err))
	}

	return nil, types.NewListTasksOutput(page), nil
}

func (r *Resolver) GetTaskTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetTaskInput) (*mcp.CallToolResult, types.GetTaskOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionTaskGet)

	prb := r.ProboService(ctx, input.ID)

	task, err := prb.Tasks.Get(ctx, input.ID)
	if err != nil {
		return nil, types.GetTaskOutput{}, fmt.Errorf("failed to get task: %w", err)
	}
	return nil, types.GetTaskOutput{
		Task: types.NewTask(task),
	}, nil
}

func (r *Resolver) AddTaskTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddTaskInput) (*mcp.CallToolResult, types.AddTaskOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionTaskCreate)

	svc := r.ProboService(ctx, input.OrganizationID)

	task, err := svc.Tasks.Create(
		ctx,
		probo.CreateTaskRequest{
			OrganizationID: input.OrganizationID,
			MeasureID:      input.MeasureID,
			Name:           input.Name,
			Description:    input.Description,
			TimeEstimate:   input.TimeEstimate,
			Deadline:       input.Deadline,
			AssignedToID:   input.AssignedToID,
		},
	)
	if err != nil {
		return nil, types.AddTaskOutput{}, fmt.Errorf("failed to create task: %w", err)
	}
	return nil, types.AddTaskOutput{
		Task: types.NewTask(task),
	}, nil
}

func (r *Resolver) UpdateTaskTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateTaskInput) (*mcp.CallToolResult, types.UpdateTaskOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionTaskUpdate)

	svc := r.ProboService(ctx, input.ID)

	task, err := svc.Tasks.Update(
		ctx,
		probo.UpdateTaskRequest{
			TaskID:       input.ID,
			Name:         input.Name,
			Description:  UnwrapOmittable(input.Description),
			State:        input.State,
			TimeEstimate: UnwrapOmittable(input.TimeEstimate),
			Deadline:     UnwrapOmittable(input.Deadline),
			AssignedToID: UnwrapOmittable(input.AssignedToID),
			MeasureID:    UnwrapOmittable(input.MeasureID),
		},
	)
	if err != nil {
		return nil, types.UpdateTaskOutput{}, fmt.Errorf("failed to update task: %w", err)
	}
	return nil, types.UpdateTaskOutput{
		Task: types.NewTask(task),
	}, nil
}

func (r *Resolver) AssignTaskTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AssignTaskInput) (*mcp.CallToolResult, types.AssignTaskOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionTaskAssign)

	svc := r.ProboService(ctx, input.ID)

	task, err := svc.Tasks.Assign(ctx, input.ID, input.AssignedToID)
	if err != nil {
		return nil, types.AssignTaskOutput{}, fmt.Errorf("failed to assign task: %w", err)
	}

	return nil, types.AssignTaskOutput{
		Task: types.NewTask(task),
	}, nil
}

func (r *Resolver) UnassignTaskTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UnassignTaskInput) (*mcp.CallToolResult, types.UnassignTaskOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionTaskUnassign)

	svc := r.ProboService(ctx, input.ID)

	task, err := svc.Tasks.Unassign(ctx, input.ID)
	if err != nil {
		return nil, types.UnassignTaskOutput{}, fmt.Errorf("failed to unassign task: %w", err)
	}
	return nil, types.UnassignTaskOutput{
		Task: types.NewTask(task),
	}, nil
}

func (r *Resolver) DeleteTaskTool(ctx context.Context, req *mcp.CallToolRequest, input *types.DeleteTaskInput) (*mcp.CallToolResult, types.DeleteTaskOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionTaskDelete)

	svc := r.ProboService(ctx, input.ID)

	err := svc.Tasks.Delete(ctx, input.ID)
	if err != nil {
		return nil, types.DeleteTaskOutput{}, fmt.Errorf("failed to delete task: %w", err)
	}

	return nil, types.DeleteTaskOutput{
		DeletedTaskID: input.ID,
	}, nil
}

func (r *Resolver) ListSnapshotsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListSnapshotsInput) (*mcp.CallToolResult, types.ListSnapshotsOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionSnapshotList)

	prb := r.ProboService(ctx, input.OrganizationID)

	pageOrderBy := page.OrderBy[coredata.SnapshotOrderField]{
		Field:     coredata.SnapshotOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.SnapshotOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	page, err := prb.Snapshots.ListForOrganizationID(ctx, input.OrganizationID, cursor)
	if err != nil {
		panic(fmt.Errorf("cannot list organization snapshots: %w", err))
	}

	return nil, types.NewListSnapshotsOutput(page), nil
}

func (r *Resolver) GetSnapshotTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetSnapshotInput) (*mcp.CallToolResult, types.GetSnapshotOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionSnapshotGet)

	prb := r.ProboService(ctx, input.ID)

	snapshot, err := prb.Snapshots.Get(ctx, input.ID)
	if err != nil {
		return nil, types.GetSnapshotOutput{}, fmt.Errorf("failed to get snapshot: %w", err)
	}
	return nil, types.GetSnapshotOutput{
		Snapshot: types.NewSnapshot(snapshot),
	}, nil
}

func (r *Resolver) TakeSnapshotTool(ctx context.Context, req *mcp.CallToolRequest, input *types.TakeSnapshotInput) (*mcp.CallToolResult, types.TakeSnapshotOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionSnapshotCreate)

	prb := r.ProboService(ctx, input.OrganizationID)

	snapshot, err := prb.Snapshots.Create(
		ctx,
		&probo.CreateSnapshotRequest{
			OrganizationID: input.OrganizationID,
			Name:           input.Name,
			Description:    input.Description,
			Type:           input.Type,
		},
	)
	if err != nil {
		return nil, types.TakeSnapshotOutput{}, fmt.Errorf("failed to take snapshot: %w", err)
	}
	return nil, types.TakeSnapshotOutput{
		Snapshot: types.NewSnapshot(snapshot),
	}, nil
}

func (r *Resolver) ListDocumentsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListDocumentsInput) (*mcp.CallToolResult, types.ListDocumentsOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionDocumentList)

	prb := r.ProboService(ctx, input.OrganizationID)

	pageOrderBy := page.OrderBy[coredata.DocumentOrderField]{
		Field:     coredata.DocumentOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.DocumentOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	documentFilter := coredata.NewDocumentFilter(nil)
	if input.Filter != nil {
		var query *string
		if input.Filter.Query != nil && *input.Filter.Query != "" {
			query = input.Filter.Query
		}

		documentFilter = coredata.NewDocumentFilter(query)
	}

	docPage, err := prb.Documents.ListByOrganizationID(ctx, input.OrganizationID, cursor, documentFilter)
	if err != nil {
		panic(fmt.Errorf("cannot list organization documents: %w", err))
	}

	approverIDsMap := make(map[gid.GID][]gid.GID)
	for _, d := range docPage.Data {
		approverPage, err := prb.Documents.ListApprovers(ctx, d.ID, allApproversCursor())
		if err != nil {
			panic(fmt.Errorf("cannot list document approvers: %w", err))
		}
		approverIDsMap[d.ID] = profileIDs(approverPage)
	}

	return nil, types.NewListDocumentsOutput(docPage, approverIDsMap), nil
}

func (r *Resolver) GetDocumentTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetDocumentInput) (*mcp.CallToolResult, types.GetDocumentOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionDocumentGet)

	prb := r.ProboService(ctx, input.ID)

	document, err := prb.Documents.Get(ctx, input.ID)
	if err != nil {
		panic(fmt.Errorf("cannot get document: %w", err))
	}

	approverPage, err := prb.Documents.ListApprovers(ctx, input.ID, allApproversCursor())
	if err != nil {
		panic(fmt.Errorf("cannot list document approvers: %w", err))
	}

	return nil, types.GetDocumentOutput{
		Document: types.NewDocument(document, profileIDs(approverPage)),
	}, nil
}

func (r *Resolver) AddDocumentTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddDocumentInput) (*mcp.CallToolResult, types.AddDocumentOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionDocumentCreate)

	svc := r.ProboService(ctx, input.OrganizationID)

	var trustCenterVisibility *coredata.TrustCenterVisibility
	if input.TrustCenterVisibility != nil {
		trustCenterVisibility = input.TrustCenterVisibility
	}

	document, documentVersion, err := svc.Documents.Create(
		ctx,
		probo.CreateDocumentRequest{
			OrganizationID:        input.OrganizationID,
			Title:                 input.Title,
			Content:               input.Content,
			ApproverIDs:           input.ApproverIds,
			Classification:        input.Classification,
			DocumentType:          input.DocumentType,
			TrustCenterVisibility: trustCenterVisibility,
		},
	)
	if err != nil {
		panic(fmt.Errorf("cannot create document: %w", err))
	}

	return nil, types.NewAddDocumentOutput(document, documentVersion, input.ApproverIds, input.ApproverIds), nil
}

func (r *Resolver) UpdateDocumentTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateDocumentInput) (*mcp.CallToolResult, types.UpdateDocumentOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionDocumentUpdate)

	svc := r.ProboService(ctx, input.ID)

	document, err := svc.Documents.Update(
		ctx,
		probo.UpdateDocumentRequest{
			DocumentID:            input.ID,
			Title:                 input.Title,
			ApproverIDs:           input.ApproverIds,
			Classification:        input.Classification,
			DocumentType:          input.DocumentType,
			TrustCenterVisibility: input.TrustCenterVisibility,
		},
	)
	if err != nil {
		panic(fmt.Errorf("cannot update document: %w", err))
	}

	approverPage, err := svc.Documents.ListApprovers(ctx, input.ID, allApproversCursor())
	if err != nil {
		panic(fmt.Errorf("cannot list document approvers: %w", err))
	}

	return nil, types.UpdateDocumentOutput{
		Document: types.NewDocument(document, profileIDs(approverPage)),
	}, nil
}

func (r *Resolver) ListDocumentVersionsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListDocumentVersionsInput) (*mcp.CallToolResult, types.ListDocumentVersionsOutput, error) {
	r.MustAuthorize(ctx, input.DocumentID, probo.ActionDocumentVersionList)

	pageOrderBy := page.OrderBy[coredata.DocumentVersionOrderField]{
		Field:     coredata.DocumentVersionOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.DocumentVersionOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)
	svc := r.ProboService(ctx, input.DocumentID)

	versionPage, err := svc.Documents.ListVersions(ctx, input.DocumentID, cursor, coredata.NewDocumentVersionFilter())
	if err != nil {
		panic(fmt.Errorf("cannot list document versions: %w", err))
	}

	approverIDsMap := make(map[gid.GID][]gid.GID)
	for _, v := range versionPage.Data {
		approverPage, err := svc.Documents.ListVersionApprovers(ctx, v.ID, allApproversCursor())
		if err != nil {
			panic(fmt.Errorf("cannot list document version approvers: %w", err))
		}
		approverIDsMap[v.ID] = profileIDs(approverPage)
	}

	return nil, types.NewListDocumentVersionsOutput(versionPage, approverIDsMap), nil
}

func (r *Resolver) GetDocumentVersionTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetDocumentVersionInput) (*mcp.CallToolResult, types.GetDocumentVersionOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionDocumentVersionGet)

	svc := r.ProboService(ctx, input.ID)

	version, err := svc.Documents.GetVersion(ctx, input.ID)
	if err != nil {
		panic(fmt.Errorf("cannot get document version: %w", err))
	}

	approverPage, err := svc.Documents.ListVersionApprovers(ctx, input.ID, allApproversCursor())
	if err != nil {
		panic(fmt.Errorf("cannot list document version approvers: %w", err))
	}

	return nil, types.GetDocumentVersionOutput{
		DocumentVersion: types.NewDocumentVersion(version, profileIDs(approverPage)),
	}, nil
}

func (r *Resolver) CreateDraftDocumentVersionTool(ctx context.Context, req *mcp.CallToolRequest, input *types.CreateDraftDocumentVersionInput) (*mcp.CallToolResult, types.CreateDraftDocumentVersionOutput, error) {
	r.MustAuthorize(ctx, input.DocumentID, probo.ActionDocumentDraftVersionCreate)

	svc := r.ProboService(ctx, input.DocumentID)

	draftVersion, err := svc.Documents.CreateDraft(ctx, input.DocumentID)
	if err != nil {
		panic(fmt.Errorf("cannot create draft document version: %w", err))
	}

	approverPage, err := svc.Documents.ListVersionApprovers(ctx, draftVersion.ID, allApproversCursor())
	if err != nil {
		panic(fmt.Errorf("cannot list document version approvers: %w", err))
	}

	return nil, types.CreateDraftDocumentVersionOutput{
		DocumentVersion: types.NewDocumentVersion(draftVersion, profileIDs(approverPage)),
	}, nil
}

func (r *Resolver) UpdateDocumentVersionTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateDocumentVersionInput) (*mcp.CallToolResult, types.UpdateDocumentVersionOutput, error) {
	r.MustAuthorize(ctx, input.DocumentVersionID, probo.ActionDocumentVersionUpdate)

	svc := r.ProboService(ctx, input.DocumentVersionID)

	documentVersion, err := svc.Documents.UpdateVersion(
		ctx,
		probo.UpdateDocumentVersionRequest{
			ID:      input.DocumentVersionID,
			Content: input.Content,
		},
	)
	if err != nil {
		panic(fmt.Errorf("cannot update document version: %w", err))
	}

	versionApproverPage, err := svc.Documents.ListVersionApprovers(ctx, documentVersion.ID, allApproversCursor())
	if err != nil {
		panic(fmt.Errorf("cannot list document version approvers: %w", err))
	}

	return nil, types.UpdateDocumentVersionOutput{
		DocumentVersion: types.NewDocumentVersion(documentVersion, profileIDs(versionApproverPage)),
	}, nil
}

func (r *Resolver) PublishDocumentVersionTool(ctx context.Context, req *mcp.CallToolRequest, input *types.PublishDocumentVersionInput) (*mcp.CallToolResult, types.PublishDocumentVersionOutput, error) {
	r.MustAuthorize(ctx, input.DocumentID, probo.ActionDocumentVersionPublish)

	svc := r.ProboService(ctx, input.DocumentID)

	user := authn.IdentityFromContext(ctx)

	document, documentVersion, err := svc.Documents.PublishVersion(ctx, input.DocumentID, user.ID, input.Changelog)
	if err != nil {
		panic(fmt.Errorf("cannot publish document version: %w", err))
	}

	docApproverPage, err := svc.Documents.ListApprovers(ctx, document.ID, allApproversCursor())
	if err != nil {
		panic(fmt.Errorf("cannot list document approvers: %w", err))
	}

	versionApproverPage, err := svc.Documents.ListVersionApprovers(ctx, documentVersion.ID, allApproversCursor())
	if err != nil {
		panic(fmt.Errorf("cannot list document version approvers: %w", err))
	}

	return nil, types.PublishDocumentVersionOutput{
		Document:        types.NewDocument(document, profileIDs(docApproverPage)),
		DocumentVersion: types.NewDocumentVersion(documentVersion, profileIDs(versionApproverPage)),
	}, nil
}

func (r *Resolver) ListDocumentVersionSignaturesTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListDocumentVersionSignaturesInput) (*mcp.CallToolResult, types.ListDocumentVersionSignaturesOutput, error) {
	r.MustAuthorize(ctx, input.DocumentVersionID, probo.ActionDocumentVersionSignatureList)

	prb := r.ProboService(ctx, input.DocumentVersionID)

	pageOrderBy := page.OrderBy[coredata.DocumentVersionSignatureOrderField]{
		Field:     coredata.DocumentVersionSignatureOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.DocumentVersionSignatureOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	var signatureStates []coredata.DocumentVersionSignatureState
	var activeContract *bool
	if input.Filter != nil {
		if input.Filter.States != nil {
			signatureStates = input.Filter.States
		}
		if input.Filter.ActiveContract != nil {
			activeContract = input.Filter.ActiveContract
		}
	}
	signatureFilter := coredata.NewDocumentVersionSignatureFilter(signatureStates, activeContract)

	page, err := prb.Documents.ListSignatures(ctx, input.DocumentVersionID, cursor, signatureFilter)
	if err != nil {
		panic(fmt.Errorf("cannot list document version signatures: %w", err))
	}

	return nil, types.NewListDocumentVersionSignaturesOutput(page), nil
}

func (r *Resolver) GetDocumentVersionSignatureTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetDocumentVersionSignatureInput) (*mcp.CallToolResult, types.GetDocumentVersionSignatureOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionDocumentVersionSignatureGet)

	prb := r.ProboService(ctx, input.ID)

	signature, err := prb.Documents.GetVersionSignature(ctx, input.ID)
	if err != nil {
		panic(fmt.Errorf("cannot get document version signature: %w", err))
	}

	return nil, types.GetDocumentVersionSignatureOutput{
		DocumentVersionSignature: types.NewDocumentVersionSignature(signature),
	}, nil
}

func (r *Resolver) RequestDocumentVersionSignatureTool(ctx context.Context, req *mcp.CallToolRequest, input *types.RequestDocumentVersionSignatureInput) (*mcp.CallToolResult, types.RequestDocumentVersionSignatureOutput, error) {
	r.MustAuthorize(ctx, input.DocumentVersionID, probo.ActionDocumentVersionSignatureRequest)

	svc := r.ProboService(ctx, input.DocumentVersionID)

	documentVersionSignature, err := svc.Documents.RequestSignature(
		ctx,
		probo.RequestSignatureRequest{
			DocumentVersionID: input.DocumentVersionID,
			Signatory:         input.SignatoryID,
		},
	)
	if err != nil {
		panic(fmt.Errorf("cannot request signature: %w", err))
	}

	return nil, types.RequestDocumentVersionSignatureOutput{
		DocumentVersionSignature: types.NewDocumentVersionSignature(documentVersionSignature),
	}, nil
}

func (r *Resolver) DeleteDraftDocumentVersionTool(ctx context.Context, req *mcp.CallToolRequest, input *types.DeleteDraftDocumentVersionInput) (*mcp.CallToolResult, types.DeleteDraftDocumentVersionOutput, error) {
	r.MustAuthorize(ctx, input.DocumentVersionID, probo.ActionDocumentVersionDeleteDraft)

	svc := r.ProboService(ctx, input.DocumentVersionID)

	err := svc.Documents.DeleteDraft(ctx, input.DocumentVersionID)
	if err != nil {
		panic(fmt.Errorf("cannot delete draft document version: %w", err))
	}

	return nil, types.DeleteDraftDocumentVersionOutput{
		DeletedDocumentVersionID: input.DocumentVersionID,
	}, nil
}

func (r *Resolver) DeleteDocumentTool(ctx context.Context, req *mcp.CallToolRequest, input *types.DeleteDocumentInput) (*mcp.CallToolResult, types.DeleteDocumentOutput, error) {
	r.MustAuthorize(ctx, input.DocumentID, probo.ActionDocumentDelete)

	svc := r.ProboService(ctx, input.DocumentID)

	err := svc.Documents.SoftDelete(ctx, input.DocumentID)
	if err != nil {
		panic(fmt.Errorf("cannot soft delete document: %w", err))
	}

	return nil, types.DeleteDocumentOutput{
		DeletedDocumentID: input.DocumentID,
	}, nil
}

func (r *Resolver) CancelSignatureRequestTool(ctx context.Context, req *mcp.CallToolRequest, input *types.CancelSignatureRequestInput) (*mcp.CallToolResult, types.CancelSignatureRequestOutput, error) {
	r.MustAuthorize(ctx, input.DocumentVersionSignatureID, probo.ActionDocumentVersionCancelSignature)

	svc := r.ProboService(ctx, input.DocumentVersionSignatureID)

	err := svc.Documents.CancelSignatureRequest(ctx, input.DocumentVersionSignatureID)
	if err != nil {
		panic(fmt.Errorf("cannot cancel signature request: %w", err))
	}

	return nil, types.CancelSignatureRequestOutput{
		DeletedDocumentVersionSignatureID: input.DocumentVersionSignatureID,
	}, nil
}

func (r *Resolver) ListMeetingsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListMeetingsInput) (*mcp.CallToolResult, types.ListMeetingsOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionMeetingList)

	prb := r.ProboService(ctx, input.OrganizationID)

	pageOrderBy := page.OrderBy[coredata.MeetingOrderField]{
		Field:     coredata.MeetingOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.MeetingOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	page, err := prb.Meetings.ListForOrganizationID(ctx, input.OrganizationID, cursor)
	if err != nil {
		panic(fmt.Errorf("cannot list organization meetings: %w", err))
	}

	return nil, types.NewListMeetingsOutput(page), nil
}

func (r *Resolver) GetMeetingTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetMeetingInput) (*mcp.CallToolResult, types.GetMeetingOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionMeetingGet)

	prb := r.ProboService(ctx, input.ID)

	meeting, err := prb.Meetings.Get(ctx, input.ID)
	if err != nil {
		return nil, types.GetMeetingOutput{}, fmt.Errorf("failed to get meeting: %w", err)
	}

	return nil, types.GetMeetingOutput{
		Meeting: types.NewMeeting(meeting),
	}, nil
}

func (r *Resolver) AddMeetingTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddMeetingInput) (*mcp.CallToolResult, types.AddMeetingOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionMeetingCreate)

	svc := r.ProboService(ctx, input.OrganizationID)

	meeting, err := svc.Meetings.Create(
		ctx,
		probo.CreateMeetingRequest{
			OrganizationID: input.OrganizationID,
			Name:           input.Name,
			Date:           input.Date,
			AttendeeIDs:    input.AttendeeIds,
			Minutes:        input.Minutes,
		},
	)
	if err != nil {
		return nil, types.AddMeetingOutput{}, fmt.Errorf("failed to create meeting: %w", err)
	}

	return nil, types.AddMeetingOutput{
		Meeting: types.NewMeeting(meeting),
	}, nil
}

func (r *Resolver) UpdateMeetingTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateMeetingInput) (*mcp.CallToolResult, types.UpdateMeetingOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionMeetingUpdate)

	svc := r.ProboService(ctx, input.ID)

	meeting, err := svc.Meetings.Update(
		ctx,
		probo.UpdateMeetingRequest{
			MeetingID:   input.ID,
			Name:        input.Name,
			Date:        input.Date,
			AttendeeIDs: input.AttendeeIds,
			Minutes:     UnwrapOmittable(input.Minutes),
		},
	)
	if err != nil {
		return nil, types.UpdateMeetingOutput{}, fmt.Errorf("failed to update meeting: %w", err)
	}

	return nil, types.UpdateMeetingOutput{
		Meeting: types.NewMeeting(meeting),
	}, nil
}

func (r *Resolver) DeleteMeetingTool(ctx context.Context, req *mcp.CallToolRequest, input *types.DeleteMeetingInput) (*mcp.CallToolResult, types.DeleteMeetingOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionMeetingDelete)

	svc := r.ProboService(ctx, input.ID)

	err := svc.Meetings.Delete(ctx, input.ID)
	if err != nil {
		return nil, types.DeleteMeetingOutput{}, fmt.Errorf("failed to delete meeting: %w", err)
	}

	return nil, types.DeleteMeetingOutput{
		DeletedMeetingID: input.ID,
	}, nil
}

func (r *Resolver) DeleteRiskTool(ctx context.Context, req *mcp.CallToolRequest, input *types.DeleteRiskInput) (*mcp.CallToolResult, types.DeleteRiskOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionRiskDelete)

	svc := r.ProboService(ctx, input.ID)

	err := svc.Risks.Delete(ctx, input.ID)
	if err != nil {
		return nil, types.DeleteRiskOutput{}, fmt.Errorf("failed to delete risk: %w", err)
	}

	return nil, types.DeleteRiskOutput{
		DeletedRiskID: input.ID,
	}, nil
}

func (r *Resolver) ListMeetingAttendeesTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListMeetingAttendeesInput) (*mcp.CallToolResult, types.ListMeetingAttendeesOutput, error) {
	r.MustAuthorize(ctx, input.MeetingID, probo.ActionMeetingGet)

	svc := r.ProboService(ctx, input.MeetingID)

	attendees, err := svc.Meetings.GetAttendees(ctx, input.MeetingID)
	if err != nil {
		return nil, types.ListMeetingAttendeesOutput{}, fmt.Errorf("failed to list meeting attendees: %w", err)
	}

	profiles := make([]*types.Profile, 0, len(attendees))
	for _, a := range attendees {
		profiles = append(profiles, types.NewProfile(a))
	}

	return nil, types.ListMeetingAttendeesOutput{
		Attendees: profiles,
	}, nil
}

func (r *Resolver) DeleteMeasureTool(ctx context.Context, req *mcp.CallToolRequest, input *types.DeleteMeasureInput) (*mcp.CallToolResult, types.DeleteMeasureOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionMeasureDelete)

	svc := r.ProboService(ctx, input.ID)

	err := svc.Measures.Delete(ctx, input.ID)
	if err != nil {
		return nil, types.DeleteMeasureOutput{}, fmt.Errorf("failed to delete measure: %w", err)
	}

	return nil, types.DeleteMeasureOutput{
		DeletedMeasureID: input.ID,
	}, nil
}

func (r *Resolver) ListMeasureRisksTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListMeasureRisksInput) (*mcp.CallToolResult, types.ListMeasureRisksOutput, error) {
	r.MustAuthorize(ctx, input.MeasureID, probo.ActionMeasureGet)

	prb := r.ProboService(ctx, input.MeasureID)

	pageOrderBy := page.OrderBy[coredata.RiskOrderField]{
		Field:     coredata.RiskOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.RiskOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	riskPage, err := prb.Risks.ListForMeasureID(ctx, input.MeasureID, cursor, coredata.NewRiskFilter(nil, nil))
	if err != nil {
		return nil, types.ListMeasureRisksOutput{}, fmt.Errorf("failed to list measure risks: %w", err)
	}

	return nil, types.NewListMeasureRisksOutput(riskPage), nil
}

func (r *Resolver) ListMeasureControlsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListMeasureControlsInput) (*mcp.CallToolResult, types.ListMeasureControlsOutput, error) {
	r.MustAuthorize(ctx, input.MeasureID, probo.ActionMeasureGet)

	prb := r.ProboService(ctx, input.MeasureID)

	pageOrderBy := page.OrderBy[coredata.ControlOrderField]{
		Field:     coredata.ControlOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.ControlOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	controlPage, err := prb.Controls.ListForMeasureID(ctx, input.MeasureID, cursor, coredata.NewControlFilter(nil))
	if err != nil {
		return nil, types.ListMeasureControlsOutput{}, fmt.Errorf("failed to list measure controls: %w", err)
	}

	return nil, types.NewListMeasureControlsOutput(controlPage), nil
}

func (r *Resolver) ListMeasureTasksTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListMeasureTasksInput) (*mcp.CallToolResult, types.ListMeasureTasksOutput, error) {
	r.MustAuthorize(ctx, input.MeasureID, probo.ActionMeasureGet)

	prb := r.ProboService(ctx, input.MeasureID)

	pageOrderBy := page.OrderBy[coredata.TaskOrderField]{
		Field:     coredata.TaskOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.TaskOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	taskPage, err := prb.Tasks.ListForMeasureID(ctx, input.MeasureID, cursor)
	if err != nil {
		return nil, types.ListMeasureTasksOutput{}, fmt.Errorf("failed to list measure tasks: %w", err)
	}

	return nil, types.NewListMeasureTasksOutput(taskPage), nil
}

func (r *Resolver) ListMeasureEvidencesTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListMeasureEvidencesInput) (*mcp.CallToolResult, types.ListMeasureEvidencesOutput, error) {
	r.MustAuthorize(ctx, input.MeasureID, probo.ActionMeasureGet)

	prb := r.ProboService(ctx, input.MeasureID)

	pageOrderBy := page.OrderBy[coredata.EvidenceOrderField]{
		Field:     coredata.EvidenceOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	evidencePage, err := prb.Evidences.ListForMeasureID(ctx, input.MeasureID, cursor)
	if err != nil {
		return nil, types.ListMeasureEvidencesOutput{}, fmt.Errorf("failed to list measure evidences: %w", err)
	}

	return nil, types.NewListMeasureEvidencesOutput(evidencePage), nil
}

func (r *Resolver) LinkMeasureTool(ctx context.Context, req *mcp.CallToolRequest, input *types.LinkMeasureInput) (*mcp.CallToolResult, types.LinkMeasureOutput, error) {
	svc := r.ProboService(ctx, input.MeasureID)

	switch input.ResourceID.EntityType() {
	case coredata.ControlEntityType:
		r.MustAuthorize(ctx, input.MeasureID, probo.ActionControlMeasureMappingCreate)
		if _, _, err := svc.Controls.CreateMeasureMapping(ctx, input.ResourceID, input.MeasureID); err != nil {
			return nil, types.LinkMeasureOutput{}, fmt.Errorf("failed to link measure to control: %w", err)
		}
	case coredata.RiskEntityType:
		r.MustAuthorize(ctx, input.MeasureID, probo.ActionRiskMeasureMappingCreate)
		if _, _, err := svc.Risks.CreateMeasureMapping(ctx, input.ResourceID, input.MeasureID); err != nil {
			return nil, types.LinkMeasureOutput{}, fmt.Errorf("failed to link measure to risk: %w", err)
		}
	default:
		return nil, types.LinkMeasureOutput{}, fmt.Errorf("unsupported resource type for measure linking: entity type %d", input.ResourceID.EntityType())
	}

	return nil, types.LinkMeasureOutput{}, nil
}

func (r *Resolver) UnlinkMeasureTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UnlinkMeasureInput) (*mcp.CallToolResult, types.UnlinkMeasureOutput, error) {
	svc := r.ProboService(ctx, input.MeasureID)

	switch input.ResourceID.EntityType() {
	case coredata.ControlEntityType:
		r.MustAuthorize(ctx, input.MeasureID, probo.ActionControlMeasureMappingDelete)
		if _, _, err := svc.Controls.DeleteMeasureMapping(ctx, input.ResourceID, input.MeasureID); err != nil {
			return nil, types.UnlinkMeasureOutput{}, fmt.Errorf("failed to unlink measure from control: %w", err)
		}
	case coredata.RiskEntityType:
		r.MustAuthorize(ctx, input.MeasureID, probo.ActionRiskMeasureMappingDelete)
		if _, _, err := svc.Risks.DeleteMeasureMapping(ctx, input.ResourceID, input.MeasureID); err != nil {
			return nil, types.UnlinkMeasureOutput{}, fmt.Errorf("failed to unlink measure from risk: %w", err)
		}
	default:
		return nil, types.UnlinkMeasureOutput{}, fmt.Errorf("unsupported resource type for measure unlinking: entity type %d", input.ResourceID.EntityType())
	}

	return nil, types.UnlinkMeasureOutput{}, nil
}

func (r *Resolver) ListUsersTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListUsersInput) (*mcp.CallToolResult, types.ListUsersOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, iam.ActionMembershipProfileList)

	pageOrderBy := page.OrderBy[coredata.MembershipProfileOrderField]{
		Field:     coredata.MembershipProfileOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.MembershipProfileOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}
	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	filter := coredata.NewMembershipProfileFilter(nil).WithMembership()
	if input.Filter != nil {
		filter = coredata.NewMembershipProfileFilter(input.Filter.ExcludeContractEnded).WithMembership()
	}

	pageResult, err := r.iamSvc.OrganizationService.ListProfiles(ctx, input.OrganizationID, cursor, filter)
	if err != nil {
		return nil, types.ListUsersOutput{}, fmt.Errorf("list users: %w", err)
	}

	users := make([]*types.Profile, 0, len(pageResult.Data))
	for _, p := range pageResult.Data {
		users = append(users, types.NewProfile(p))
	}
	var nextCursor *page.CursorKey
	if len(pageResult.Data) > 0 && pageResult.Cursor != nil {
		cursorKey := pageResult.Data[len(pageResult.Data)-1].CursorKey(pageResult.Cursor.OrderBy.Field)
		nextCursor = &cursorKey
	}
	return nil, types.ListUsersOutput{
		Users:      users,
		NextCursor: nextCursor,
	}, nil
}

func (r *Resolver) GetUserTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetUserInput) (*mcp.CallToolResult, types.GetUserOutput, error) {
	profile, err := r.iamSvc.OrganizationService.GetProfile(ctx, input.ID)
	if err != nil {
		var errNotFound *iam.ErrProfileNotFound
		if errors.As(err, &errNotFound) {
			return nil, types.GetUserOutput{}, fmt.Errorf("user not found: %w", err)
		}
		return nil, types.GetUserOutput{}, fmt.Errorf("get user: %w", err)
	}
	r.MustAuthorize(ctx, profile.OrganizationID, iam.ActionMembershipProfileGet)
	return nil, types.GetUserOutput{User: types.NewProfile(profile)}, nil
}

func (r *Resolver) CreateUserTool(ctx context.Context, req *mcp.CallToolRequest, input *types.CreateUserInput) (*mcp.CallToolResult, types.CreateUserOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, iam.ActionMembershipProfileCreate)

	var contractStart, contractEnd **time.Time
	if input.ContractStartDate != nil {
		contractStart = &input.ContractStartDate
	}
	if input.ContractEndDate != nil {
		contractEnd = &input.ContractEndDate
	}
	profile, err := r.iamSvc.OrganizationService.CreateUser(ctx, &iam.CreateUserRequest{
		OrganizationID:           input.OrganizationID,
		EmailAddress:             input.EmailAddress,
		Role:                     input.Role,
		FullName:                 input.FullName,
		AdditionalEmailAddresses: input.AdditionalEmailAddresses,
		Kind:                     input.Kind,
		Position:                 input.Position,
		ContractStartDate:        contractStart,
		ContractEndDate:          contractEnd,
	})
	if err != nil {
		var errAlreadyExists *iam.ErrUserAlreadyExists
		if errors.As(err, &errAlreadyExists) {
			return nil, types.CreateUserOutput{}, fmt.Errorf("user with email already exists: %w", err)
		}
		return nil, types.CreateUserOutput{}, fmt.Errorf("create user: %w", err)
	}
	return nil, types.CreateUserOutput{User: types.NewProfile(profile)}, nil
}

func (r *Resolver) InviteUserTool(ctx context.Context, req *mcp.CallToolRequest, input *types.InviteUserInput) (*mcp.CallToolResult, types.InviteUserOutput, error) {
	r.MustAuthorize(ctx, input.ProfileID, iam.ActionInvitationCreate)

	invitation, err := r.iamSvc.OrganizationService.InviteUser(ctx, &iam.CreateInvitationRequest{
		OrganizationID: input.OrganizationID,
		ProfileID:      input.ProfileID,
	})
	if err != nil {
		var errOrgNotFound *iam.ErrOrganizationNotFound
		var errUserExists *iam.ErrUserAlreadyExists
		if errors.As(err, &errOrgNotFound) {
			return nil, types.InviteUserOutput{}, fmt.Errorf("organization not found: %w", err)
		}
		if errors.As(err, &errUserExists) {
			return nil, types.InviteUserOutput{}, fmt.Errorf("user already in organization: %w", err)
		}
		return nil, types.InviteUserOutput{}, fmt.Errorf("invite user: %w", err)
	}
	return nil, types.InviteUserOutput{InvitationID: invitation.ID}, nil
}

func (r *Resolver) UpdateUserTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateUserInput) (*mcp.CallToolResult, types.UpdateUserOutput, error) {
	r.MustAuthorize(ctx, input.ID, iam.ActionMembershipProfileUpdate)

	var additionalEmails []mail.Addr
	if input.AdditionalEmailAddresses != nil {
		additionalEmails = *input.AdditionalEmailAddresses
	}
	var position *string
	if p := UnwrapOmittable(input.Position); p != nil {
		position = *p
	}
	var contractStart, contractEnd **time.Time
	if p := UnwrapOmittable(input.ContractStartDate); p != nil {
		contractStart = p
	}
	if p := UnwrapOmittable(input.ContractEndDate); p != nil {
		contractEnd = p
	}
	profile, err := r.iamSvc.OrganizationService.UpdateUser(ctx, &iam.UpdateUserRequest{
		ID:                       input.ID,
		FullName:                 input.FullName,
		AdditionalEmailAddresses: additionalEmails,
		Kind:                     input.Kind,
		Position:                 position,
		ContractStartDate:        contractStart,
		ContractEndDate:          contractEnd,
	})
	if err != nil {
		return nil, types.UpdateUserOutput{}, fmt.Errorf("update user: %w", err)
	}
	return nil, types.UpdateUserOutput{User: types.NewProfile(profile)}, nil
}

func (r *Resolver) UpdateMembershipTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateMembershipInput) (*mcp.CallToolResult, types.UpdateMembershipOutput, error) {
	r.MustAuthorize(ctx, input.MembershipID, iam.ActionMembershipUpdate)
	if input.Role == coredata.MembershipRoleOwner {
		r.MustAuthorize(ctx, input.MembershipID, iam.ActionMembershipRoleSetOwner)
	}

	membership, err := r.iamSvc.OrganizationService.UpdateMempership(ctx, input.OrganizationID, input.MembershipID, input.Role)
	if err != nil {
		return nil, types.UpdateMembershipOutput{}, fmt.Errorf("update membership: %w", err)
	}
	return nil, types.UpdateMembershipOutput{
		Membership: &types.Membership{
			ID:        membership.ID,
			Role:      membership.Role,
			CreatedAt: membership.CreatedAt,
		},
	}, nil
}

func (r *Resolver) RemoveUserTool(ctx context.Context, req *mcp.CallToolRequest, input *types.RemoveUserInput) (*mcp.CallToolResult, types.RemoveUserOutput, error) {
	r.MustAuthorize(ctx, input.ProfileID, iam.ActionMembershipProfileDelete)

	err := r.iamSvc.OrganizationService.RemoveUser(ctx, input.OrganizationID, input.ProfileID)
	if err != nil {
		var errManagedBySCIM *iam.ErrUserManagedBySCIM
		var errLastOwner *iam.ErrLastActiveOwner
		if errors.As(err, &errManagedBySCIM) {
			return nil, types.RemoveUserOutput{}, fmt.Errorf("user is managed by SCIM and cannot be removed: %w", err)
		}
		if errors.As(err, &errLastOwner) {
			return nil, types.RemoveUserOutput{}, fmt.Errorf("cannot remove last active owner: %w", err)
		}
		return nil, types.RemoveUserOutput{}, fmt.Errorf("remove user: %w", err)
	}
	return nil, types.RemoveUserOutput{DeletedUserID: input.ProfileID}, nil
}

func (r *Resolver) DeleteDataProtectionImpactAssessmentTool(ctx context.Context, req *mcp.CallToolRequest, input *types.DeleteDataProtectionImpactAssessmentInput) (*mcp.CallToolResult, types.DeleteDataProtectionImpactAssessmentOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionDataProtectionImpactAssessmentDelete)

	svc := r.ProboService(ctx, input.ID)

	err := svc.DataProtectionImpactAssessments.Delete(ctx, input.ID)
	if err != nil {
		return nil, types.DeleteDataProtectionImpactAssessmentOutput{}, fmt.Errorf("failed to delete data protection impact assessment: %w", err)
	}

	return nil, types.DeleteDataProtectionImpactAssessmentOutput{
		DeletedDataProtectionImpactAssessmentID: input.ID,
	}, nil
}

func (r *Resolver) ListStatesOfApplicabilityTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListStatesOfApplicabilityInput) (*mcp.CallToolResult, types.ListStatesOfApplicabilityOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionStateOfApplicabilityList)

	prb := r.ProboService(ctx, input.OrganizationID)

	pageOrderBy := page.OrderBy[coredata.StateOfApplicabilityOrderField]{
		Field:     coredata.StateOfApplicabilityOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.StateOfApplicabilityOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	pg, err := prb.StatesOfApplicability.ListForOrganizationID(ctx, input.OrganizationID, cursor, coredata.NewStateOfApplicabilityFilter(nil))
	if err != nil {
		return nil, types.ListStatesOfApplicabilityOutput{}, fmt.Errorf("failed to list states of applicability: %w", err)
	}

	return nil, types.NewListStatesOfApplicabilityOutput(pg), nil
}

func (r *Resolver) GetStateOfApplicabilityTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetStateOfApplicabilityInput) (*mcp.CallToolResult, types.GetStateOfApplicabilityOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionStateOfApplicabilityGet)

	prb := r.ProboService(ctx, input.ID)

	soa, err := prb.StatesOfApplicability.Get(ctx, input.ID)
	if err != nil {
		return nil, types.GetStateOfApplicabilityOutput{}, fmt.Errorf("failed to get state of applicability: %w", err)
	}

	return nil, types.GetStateOfApplicabilityOutput{
		StateOfApplicability: types.NewStateOfApplicability(soa),
	}, nil
}

func (r *Resolver) AddStateOfApplicabilityTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddStateOfApplicabilityInput) (*mcp.CallToolResult, types.AddStateOfApplicabilityOutput, error) {
	r.MustAuthorize(ctx, input.OrganizationID, probo.ActionStateOfApplicabilityCreate)

	svc := r.ProboService(ctx, input.OrganizationID)

	soa, err := svc.StatesOfApplicability.Create(ctx, probo.CreateStateOfApplicabilityRequest{
		OrganizationID: input.OrganizationID,
		Name:           input.Name,
		OwnerID:        input.OwnerID,
	})
	if err != nil {
		return nil, types.AddStateOfApplicabilityOutput{}, fmt.Errorf("failed to create state of applicability: %w", err)
	}

	return nil, types.AddStateOfApplicabilityOutput{
		StateOfApplicability: types.NewStateOfApplicability(soa),
	}, nil
}

func (r *Resolver) UpdateStateOfApplicabilityTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateStateOfApplicabilityInput) (*mcp.CallToolResult, types.UpdateStateOfApplicabilityOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionStateOfApplicabilityUpdate)

	svc := r.ProboService(ctx, input.ID)

	soa, err := svc.StatesOfApplicability.Update(ctx, probo.UpdateStateOfApplicabilityRequest{
		StateOfApplicabilityID: input.ID,
		Name:                   input.Name,
		OwnerID:                input.OwnerID,
	})
	if err != nil {
		return nil, types.UpdateStateOfApplicabilityOutput{}, fmt.Errorf("failed to update state of applicability: %w", err)
	}

	return nil, types.UpdateStateOfApplicabilityOutput{
		StateOfApplicability: types.NewStateOfApplicability(soa),
	}, nil
}

func (r *Resolver) DeleteStateOfApplicabilityTool(ctx context.Context, req *mcp.CallToolRequest, input *types.DeleteStateOfApplicabilityInput) (*mcp.CallToolResult, types.DeleteStateOfApplicabilityOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionStateOfApplicabilityDelete)

	svc := r.ProboService(ctx, input.ID)

	err := svc.StatesOfApplicability.Delete(ctx, input.ID)
	if err != nil {
		return nil, types.DeleteStateOfApplicabilityOutput{}, fmt.Errorf("failed to delete state of applicability: %w", err)
	}

	return nil, types.DeleteStateOfApplicabilityOutput{
		DeletedStateOfApplicabilityID: input.ID,
	}, nil
}

func (r *Resolver) ExportStateOfApplicabilityPDFTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ExportStateOfApplicabilityPDFInput) (*mcp.CallToolResult, types.ExportStateOfApplicabilityPDFOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionStateOfApplicabilityExport)

	svc := r.ProboService(ctx, input.ID)

	soa, err := svc.StatesOfApplicability.Get(ctx, input.ID)
	if err != nil {
		return nil, types.ExportStateOfApplicabilityPDFOutput{}, fmt.Errorf("failed to get state of applicability: %w", err)
	}

	pdfData, err := svc.StatesOfApplicability.ExportPDF(ctx, input.ID)
	if err != nil {
		return nil, types.ExportStateOfApplicabilityPDFOutput{}, fmt.Errorf("failed to export state of applicability PDF: %w", err)
	}

	return nil, types.ExportStateOfApplicabilityPDFOutput{
		PdfBase64: base64.StdEncoding.EncodeToString(pdfData),
		Filename:  soa.Name + ".pdf",
	}, nil
}

func (r *Resolver) ListApplicabilityStatementsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListApplicabilityStatementsInput) (*mcp.CallToolResult, types.ListApplicabilityStatementsOutput, error) {
	r.MustAuthorize(ctx, input.StateOfApplicabilityID, probo.ActionApplicabilityStatementList)

	prb := r.ProboService(ctx, input.StateOfApplicabilityID)

	pageOrderBy := page.OrderBy[coredata.ApplicabilityStatementOrderField]{
		Field:     coredata.ApplicabilityStatementOrderFieldControlSectionTitle,
		Direction: page.OrderDirectionAsc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.ApplicabilityStatementOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	pg, err := prb.StatesOfApplicability.ListApplicabilityStatements(ctx, input.StateOfApplicabilityID, cursor)
	if err != nil {
		return nil, types.ListApplicabilityStatementsOutput{}, fmt.Errorf("failed to list applicability statements: %w", err)
	}

	return nil, types.NewListApplicabilityStatementsOutput(pg), nil
}

func (r *Resolver) GetApplicabilityStatementTool(ctx context.Context, req *mcp.CallToolRequest, input *types.GetApplicabilityStatementInput) (*mcp.CallToolResult, types.GetApplicabilityStatementOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionApplicabilityStatementGet)

	prb := r.ProboService(ctx, input.ID)

	stmt, err := prb.StatesOfApplicability.GetApplicabilityStatement(ctx, input.ID)
	if err != nil {
		return nil, types.GetApplicabilityStatementOutput{}, fmt.Errorf("failed to get applicability statement: %w", err)
	}

	return nil, types.GetApplicabilityStatementOutput{
		ApplicabilityStatement: types.NewApplicabilityStatement(stmt),
	}, nil
}

func (r *Resolver) AddApplicabilityStatementTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddApplicabilityStatementInput) (*mcp.CallToolResult, types.AddApplicabilityStatementOutput, error) {
	r.MustAuthorize(ctx, input.StateOfApplicabilityID, probo.ActionApplicabilityStatementCreate)

	svc := r.ProboService(ctx, input.StateOfApplicabilityID)

	stmt, err := svc.StatesOfApplicability.CreateApplicabilityStatement(
		ctx,
		input.StateOfApplicabilityID,
		input.ControlID,
		input.Applicability,
		input.Justification,
	)
	if err != nil {
		return nil, types.AddApplicabilityStatementOutput{}, fmt.Errorf("failed to create applicability statement: %w", err)
	}

	return nil, types.AddApplicabilityStatementOutput{
		ApplicabilityStatement: types.NewApplicabilityStatement(stmt),
	}, nil
}

func (r *Resolver) UpdateApplicabilityStatementTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateApplicabilityStatementInput) (*mcp.CallToolResult, types.UpdateApplicabilityStatementOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionApplicabilityStatementUpdate)

	svc := r.ProboService(ctx, input.ID)

	stmt, err := svc.StatesOfApplicability.UpdateApplicabilityStatement(
		ctx,
		input.ID,
		input.Applicability,
		input.Justification,
	)
	if err != nil {
		return nil, types.UpdateApplicabilityStatementOutput{}, fmt.Errorf("failed to update applicability statement: %w", err)
	}

	return nil, types.UpdateApplicabilityStatementOutput{
		ApplicabilityStatement: types.NewApplicabilityStatement(stmt),
	}, nil
}

func (r *Resolver) DeleteApplicabilityStatementTool(ctx context.Context, req *mcp.CallToolRequest, input *types.DeleteApplicabilityStatementInput) (*mcp.CallToolResult, types.DeleteApplicabilityStatementOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionApplicabilityStatementDelete)

	svc := r.ProboService(ctx, input.ID)

	err := svc.StatesOfApplicability.DeleteApplicabilityStatement(ctx, input.ID)
	if err != nil {
		return nil, types.DeleteApplicabilityStatementOutput{}, fmt.Errorf("failed to delete applicability statement: %w", err)
	}

	return nil, types.DeleteApplicabilityStatementOutput{
		DeletedApplicabilityStatementID: input.ID,
	}, nil
}

// ListVendorRiskAssessmentsTool handles the listVendorRiskAssessments tool
// List all risk assessments for a vendor
func (r *Resolver) ListVendorRiskAssessmentsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListVendorRiskAssessmentsInput) (*mcp.CallToolResult, types.ListVendorRiskAssessmentsOutput, error) {
	r.MustAuthorize(ctx, input.VendorID, probo.ActionVendorRiskAssessmentList)

	prb := r.ProboService(ctx, input.VendorID)

	pageOrderBy := page.OrderBy[coredata.VendorRiskAssessmentOrderField]{
		Field:     coredata.VendorRiskAssessmentOrderFieldCreatedAt,
		Direction: page.OrderDirectionDesc,
	}
	if input.OrderBy != nil {
		pageOrderBy = page.OrderBy[coredata.VendorRiskAssessmentOrderField]{
			Field:     input.OrderBy.Field,
			Direction: input.OrderBy.Direction,
		}
	}

	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)

	p, err := prb.Vendors.ListRiskAssessments(ctx, input.VendorID, cursor)
	if err != nil {
		return nil, types.ListVendorRiskAssessmentsOutput{}, fmt.Errorf("cannot list vendor risk assessments: %w", err)
	}

	return nil, types.NewListVendorRiskAssessmentsOutput(p), nil
}

// AddVendorRiskAssessmentTool handles the addVendorRiskAssessment tool
// Add a new risk assessment for a vendor
func (r *Resolver) AddVendorRiskAssessmentTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddVendorRiskAssessmentInput) (*mcp.CallToolResult, types.AddVendorRiskAssessmentOutput, error) {
	r.MustAuthorize(ctx, input.VendorID, probo.ActionVendorRiskAssessmentCreate)

	prb := r.ProboService(ctx, input.VendorID)

	assessment, err := prb.Vendors.CreateRiskAssessment(
		ctx,
		probo.CreateVendorRiskAssessmentRequest{
			VendorID:        input.VendorID,
			ExpiresAt:       input.ExpiresAt,
			DataSensitivity: input.DataSensitivity,
			BusinessImpact:  input.BusinessImpact,
			Notes:           input.Notes,
		},
	)
	if err != nil {
		return nil, types.AddVendorRiskAssessmentOutput{}, fmt.Errorf("failed to create vendor risk assessment: %w", err)
	}

	return nil, types.NewAddVendorRiskAssessmentOutput(assessment), nil
}

func (r *Resolver) DeleteVendorTool(ctx context.Context, req *mcp.CallToolRequest, input *types.DeleteVendorInput) (*mcp.CallToolResult, types.DeleteVendorOutput, error) {
	r.MustAuthorize(ctx, input.ID, probo.ActionVendorDelete)

	svc := r.ProboService(ctx, input.ID)

	err := svc.Vendors.Delete(ctx, input.ID)
	if err != nil {
		return nil, types.DeleteVendorOutput{}, fmt.Errorf("failed to delete vendor: %w", err)
	}

	return nil, types.DeleteVendorOutput{
		DeletedVendorID: input.ID,
	}, nil
}