package tidbtools

import (
	"context"

	premiumdm "github.com/tidbcloud/dataflow-service-ng/idl/pbgen/proto/premium_dm/v1beta2"
	"github.com/tidbcloud/dataflow-service/idl/pbgen/proto/dm"
	"google.golang.org/grpc/status"
)

func (s *svc) ListDMJobs(ctx context.Context, orgID, projectID, clusterID uint64, page, pageSize int32) (*dm.ListJobsResp, error) {
	// The DM service requires a valid projectID for permission checks.
	// Return an empty list when projectID is missing to avoid a 500.
	if projectID == 0 {
		return &dm.ListJobsResp{}, nil
	}
	ctx, cancel := context.WithTimeout(ctx, rpcTimeout)
	defer cancel()
	var err error
	if ctx, err = withDataflowOpsUser(ctx, orgID, projectID); err != nil {
		return nil, err
	}
	resp, err := s.dataflow.DM.ListJobs(ctx, &dm.ListJobsReq{
		OrgId:     orgID,
		ProjectId: projectID,
		ClusterId: clusterID,
		Page:      page,
		PageSize:  pageSize,
	})
	if err != nil {
		if st, ok := status.FromError(err); ok && st.Code() == codeDMClusterNotFound {
			return &dm.ListJobsResp{}, nil
		}
		return nil, err
	}
	return resp, nil
}

func (s *svc) GetDMJob(ctx context.Context, orgID, projectID, clusterID, jobID uint64) (*dm.GetJobResp, error) {
	ctx, cancel := context.WithTimeout(ctx, rpcTimeout)
	defer cancel()
	var err error
	if ctx, err = withDataflowOpsUser(ctx, orgID, projectID); err != nil {
		return nil, err
	}
	return s.dataflow.DM.GetJob(ctx, &dm.GetJobReq{
		OrgId:     orgID,
		ProjectId: projectID,
		ClusterId: clusterID,
		Id:        jobID,
	})
}

func (s *svc) GetDMPrecheck(ctx context.Context, orgID, projectID, clusterID, precheckID uint64) (*dm.GetPrecheckJobResp, error) {
	ctx, cancel := context.WithTimeout(ctx, rpcTimeout)
	defer cancel()
	var err error
	if ctx, err = withDataflowOpsUser(ctx, orgID, projectID); err != nil {
		return nil, err
	}
	return s.dataflow.DM.GetPrecheck(ctx, &dm.GetPrecheckJobReq{
		OrgId:     orgID,
		ProjectId: projectID,
		ClusterId: clusterID,
		Id:        precheckID,
	})
}

func (s *svc) ListDMPrechecks(ctx context.Context, orgID, projectID, clusterID uint64, page, pageSize int32) (*dm.ListPrechecksResp, error) {
	ctx, cancel := context.WithTimeout(ctx, rpcTimeout)
	defer cancel()
	var err error
	if ctx, err = withDataflowOpsUser(ctx, orgID, projectID); err != nil {
		return nil, err
	}
	return s.dataflow.DM.ListPrechecks(ctx, &dm.ListPrechecksReq{
		OrgId:     orgID,
		ProjectId: projectID,
		ClusterId: clusterID,
		Page:      page,
		PageSize:  pageSize,
	})
}

// ----- Premium DM methods (dataflow-service-ng PremiumMigrationService) -----

func (s *svc) ListPremiumDMJobs(ctx context.Context, tidbID string, pageSize int32) (*premiumdm.ListPremiumMigrationsResp, error) {
	ctx, cancel := context.WithTimeout(ctx, rpcTimeout)
	defer cancel()
	return s.dataflow.PremiumDM.ListPremiumMigrations(ctx, &premiumdm.ListPremiumMigrationsReq{
		TidbId:   tidbID,
		PageSize: pageSize,
	})
}

func (s *svc) GetPremiumDMJob(ctx context.Context, tidbID, migrationID string) (*premiumdm.PremiumMigration, error) {
	ctx, cancel := context.WithTimeout(ctx, rpcTimeout)
	defer cancel()
	return s.dataflow.PremiumDM.GetPremiumMigration(ctx, &premiumdm.GetPremiumMigrationReq{
		TidbId:      tidbID,
		MigrationId: migrationID,
	})
}

func (s *svc) ListPremiumDMPrechecks(ctx context.Context, tidbID string, pageSize int32) (*premiumdm.ListPremiumMigrationPrechecksResp, error) {
	ctx, cancel := context.WithTimeout(ctx, rpcTimeout)
	defer cancel()
	return s.dataflow.PremiumDM.ListPremiumMigrationPrechecks(ctx, &premiumdm.ListPremiumMigrationPrechecksReq{
		TidbId:   tidbID,
		PageSize: pageSize,
	})
}

func (s *svc) GetPremiumDMPrecheck(ctx context.Context, tidbID, precheckID string) (*premiumdm.PremiumMigrationPrecheck, error) {
	ctx, cancel := context.WithTimeout(ctx, rpcTimeout)
	defer cancel()
	return s.dataflow.PremiumDM.PremiumMigrationGetPrecheck(ctx, &premiumdm.GetPremiumMigrationPrecheckReq{
		TidbId:     tidbID,
		PrecheckId: precheckID,
	})
}