package agent

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/pingcap/errors"
	pkgcrypto "github.com/tidbcloud/dataflow-service-ng/pkg/crypto"
	agentpb "github.com/tidbcloud/dataflow-service/dataflow-agent-ng/idl/pbgen/proto/ng_agent"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"
)

func (c *NGRegionalClient) GetImportFiles(
	ctx context.Context, clusterID uint64, sourceURI string, maxScanFiles int32, tableMetas ...*agentpb.GetImportFilesReq_TableMeta,
) (*agentpb.GetImportFilesResp, error) {
	req := &agentpb.GetImportFilesReq{ClusterId: clusterID, SourceUri: sourceURI, TableMeta: tableMetas}
	if maxScanFiles > 0 {
		req.MaxScanFiles = maxScanFiles
	}
	url := fmt.Sprintf("%s/api/v1/import/get-import-files", c.AgentUrl)
	respBytes, err := c.getResponseProtoJson(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp agentpb.GetImportFilesResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) CreateImportFilesTask(
	ctx context.Context, taskID string, clusterID uint64, sourceURI string, maxScanFiles int32, tableMetas ...*agentpb.GetImportFilesReq_TableMeta,
) (*agentpb.CreateImportFilesTaskResp, error) {
	req := &agentpb.CreateImportFilesTaskReq{
		TaskId:       taskID,
		ClusterId:    clusterID,
		SourceUri:    sourceURI,
		MaxScanFiles: maxScanFiles,
		TableMeta:    tableMetas,
	}
	url := fmt.Sprintf("%s/api/v1/import/files-tasks", c.AgentUrl)
	respBytes, err := c.getResponseProtoJson(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp agentpb.CreateImportFilesTaskResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) GetImportFilesTask(ctx context.Context, taskID string) (*agentpb.GetImportFilesTaskResp, error) {
	url := fmt.Sprintf("%s/api/v1/import/files-tasks/%s", c.AgentUrl, taskID)
	respBytes, err := c.getResponseWithRetry(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	var resp agentpb.GetImportFilesTaskResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) postProtoRequest(ctx context.Context, url string, req proto.Message, resp proto.Message) error {
	respBytes, err := c.getResponseProtoJson(ctx, http.MethodPost, url, req)
	if err != nil {
		return err
	}
	if err := protojson.Unmarshal(respBytes, resp); err != nil {
		return errors.Trace(err)
	}
	return nil
}

func (c *NGRegionalClient) postProtoRequestWithCluster(
	ctx context.Context,
	url string,
	clusterID uint64,
	username string,
	password string,
	buildReq func(cluster *agentpb.Cluster) proto.Message,
	resp proto.Message,
) error {
	cluster, err := c.buildEncryptedCluster(ctx, clusterID, username, password)
	if err != nil {
		return err
	}
	return c.postProtoRequest(ctx, url, buildReq(cluster), resp)
}

func (c *NGRegionalClient) CreateImportIntoJobTask(
	ctx context.Context,
	taskID string,
	clusterID uint64,
	username string,
	password string,
	sql string,
	setupSQLs []string,
) (*agentpb.CreateImportIntoJobTaskResp, error) {
	url := fmt.Sprintf("%s/api/v1/import/import-into-job-tasks", c.AgentUrl)
	var resp agentpb.CreateImportIntoJobTaskResp
	if err := c.postProtoRequestWithCluster(ctx, url, clusterID, username, password, func(cluster *agentpb.Cluster) proto.Message {
		return &agentpb.CreateImportIntoJobTaskReq{
			TaskId:    taskID,
			ClusterId: clusterID,
			Cluster:   cluster,
			Sql:       sql,
			SetupSqls: setupSQLs,
		}
	}, &resp); err != nil {
		return nil, err
	}
	return &resp, nil
}

func (c *NGRegionalClient) BatchCreateImportIntoJobTasks(
	ctx context.Context,
	clusterID uint64,
	username string,
	password string,
	tasks []*agentpb.BatchCreateImportIntoJobTasksReq_Task,
) (*agentpb.BatchCreateImportIntoJobTasksResp, error) {
	url := fmt.Sprintf("%s/api/v1/import/import-into-job-tasks:batchCreate", c.AgentUrl)
	var resp agentpb.BatchCreateImportIntoJobTasksResp
	if err := c.postProtoRequestWithCluster(ctx, url, clusterID, username, password, func(cluster *agentpb.Cluster) proto.Message {
		return &agentpb.BatchCreateImportIntoJobTasksReq{
			ClusterId: clusterID,
			Cluster:   cluster,
			Tasks:     tasks,
		}
	}, &resp); err != nil {
		return nil, err
	}
	return &resp, nil
}

func (c *NGRegionalClient) GetImportIntoJobTask(ctx context.Context, taskID string) (*agentpb.GetImportIntoJobTaskResp, error) {
	url := fmt.Sprintf("%s/api/v1/import/import-into-job-tasks/%s", c.AgentUrl, taskID)
	respBytes, err := c.getResponseWithRetry(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	var resp agentpb.GetImportIntoJobTaskResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) BatchGetImportIntoJobTasks(ctx context.Context, taskIDs []string) (*agentpb.BatchGetImportIntoJobTasksResp, error) {
	req := &agentpb.BatchGetImportIntoJobTasksReq{TaskIds: taskIDs}
	url := fmt.Sprintf("%s/api/v1/import/import-into-job-tasks:batchGet", c.AgentUrl)
	var resp agentpb.BatchGetImportIntoJobTasksResp
	if err := c.postProtoRequest(ctx, url, req, &resp); err != nil {
		return nil, err
	}
	return &resp, nil
}

func (c *NGRegionalClient) BatchCancelImportIntoJobs(
	ctx context.Context,
	clusterID uint64,
	username string,
	password string,
	jobIDs []uint64,
) (*agentpb.BatchCancelImportIntoJobsResp, error) {
	url := fmt.Sprintf("%s/api/v1/import/import-into-jobs:batchCancel", c.AgentUrl)
	var resp agentpb.BatchCancelImportIntoJobsResp
	if err := c.postProtoRequestWithCluster(ctx, url, clusterID, username, password, func(cluster *agentpb.Cluster) proto.Message {
		return &agentpb.BatchCancelImportIntoJobsReq{
			ClusterId: clusterID,
			Cluster:   cluster,
			JobIds:    jobIDs,
		}
	}, &resp); err != nil {
		return nil, err
	}
	return &resp, nil
}

func (c *NGRegionalClient) CreateSchemasTask(
	ctx context.Context,
	taskID string,
	clusterID uint64,
	sourceURI string,
	username string,
	password string,
	schemas ...*agentpb.CreateSchemasReq_Schema,
) (*agentpb.CreateSchemasTaskResp, error) {
	cluster, err := c.buildEncryptedCluster(ctx, clusterID, username, password)
	if err != nil {
		return nil, err
	}
	req := &agentpb.CreateSchemasTaskReq{
		TaskId:    taskID,
		ClusterId: clusterID,
		SourceUri: sourceURI,
		Schemas:   schemas,
		Cluster:   cluster,
	}
	url := fmt.Sprintf("%s/api/v1/import/create-schemas-tasks", c.AgentUrl)
	respBytes, err := c.getResponseProtoJson(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp agentpb.CreateSchemasTaskResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) GetSchemasTask(ctx context.Context, taskID string) (*agentpb.GetSchemasTaskResp, error) {
	url := fmt.Sprintf("%s/api/v1/import/create-schemas-tasks/%s", c.AgentUrl, taskID)
	respBytes, err := c.getResponseWithRetry(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	var resp agentpb.GetSchemasTaskResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) CreateImportDataSizeEstimateTask(
	ctx context.Context,
	taskID string,
	clusterID uint64,
	sourceURI string,
	maxScanFiles int32,
	csvConfig *agentpb.CsvConfig,
	tableMetas ...*agentpb.CreateImportDataSizeEstimateTaskReq_TableMeta,
) (*agentpb.CreateImportDataSizeEstimateTaskResp, error) {
	req := &agentpb.CreateImportDataSizeEstimateTaskReq{
		TaskId:       taskID,
		ClusterId:    clusterID,
		SourceUri:    sourceURI,
		MaxScanFiles: maxScanFiles,
		TableMeta:    tableMetas,
		CsvConfig:    csvConfig,
	}
	url := fmt.Sprintf("%s/api/v1/import/estimate-data-size-tasks", c.AgentUrl)
	respBytes, err := c.getResponseProtoJson(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp agentpb.CreateImportDataSizeEstimateTaskResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) GetImportDataSizeEstimateTask(ctx context.Context, taskID string) (*agentpb.GetImportDataSizeEstimateTaskResp, error) {
	url := fmt.Sprintf("%s/api/v1/import/estimate-data-size-tasks/%s", c.AgentUrl, taskID)
	respBytes, err := c.getResponseWithRetry(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	var resp agentpb.GetImportDataSizeEstimateTaskResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) buildEncryptedCluster(ctx context.Context, clusterID uint64, username, password string) (*agentpb.Cluster, error) {
	pemResp, err := c.GetPublicKey(ctx, &agentpb.GetPublicKeyRequest{})
	if err != nil {
		return nil, err
	}
	data, err := pkgcrypto.EncryptWithPEM(pemResp.GetPem(), []byte(password))
	if err != nil {
		return nil, err
	}
	return &agentpb.Cluster{
		ClusterId:         clusterID,
		User:              username,
		PasswordEncrypted: true,
		EncryptedPassword: &agentpb.EncryptedData{
			Data: data,
			Pem:  pemResp.GetPem(),
		},
	}, nil
}

func (c *NGRegionalClient) CreateTiKVCapacityReservation(
	ctx context.Context,
	sharedPoolID, reservationID string,
	requiredDataBytes uint64,
) (*TiKVCapacityReservation, error) {
	payload, err := json.Marshal(map[string]any{
		"sharedPoolId":      sharedPoolID,
		"bizType":           tikvCapacityReservationBizTypeImport,
		"reservationId":     reservationID,
		"requiredDataBytes": requiredDataBytes,
	})
	if err != nil {
		return nil, errors.Trace(err)
	}

	url := fmt.Sprintf("%s/api/premium/v1/sharedPools/%s/tikvCapacityReservations", c.BaseUrl, sharedPoolID)
	respBytes, err := c.getResponseWithRetry(ctx, http.MethodPost, url, payload)
	if err != nil {
		return nil, err
	}
	return parseTiKVCapacityReservation(respBytes)
}

func (c *NGRegionalClient) GetTiKVCapacityReservation(ctx context.Context, sharedPoolID, reservationID string) (*TiKVCapacityReservation, error) {
	url := fmt.Sprintf("%s/api/premium/v1/sharedPools/%s/tikvCapacityReservations/%s/%s", c.BaseUrl, sharedPoolID, tikvCapacityReservationBizTypeImport, reservationID)
	respBytes, err := c.getResponseWithRetry(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	return parseTiKVCapacityReservation(respBytes)
}

func (c *NGRegionalClient) DeleteTiKVCapacityReservation(ctx context.Context, sharedPoolID, reservationID string) error {
	url := fmt.Sprintf("%s/api/premium/v1/sharedPools/%s/tikvCapacityReservations/%s/%s", c.BaseUrl, sharedPoolID, tikvCapacityReservationBizTypeImport, reservationID)
	_, err := c.getResponseWithRetry(ctx, http.MethodDelete, url, nil)
	return err
}

func parseTiKVCapacityReservation(respBytes []byte) (*TiKVCapacityReservation, error) {
	var resp TiKVCapacityReservation
	if err := json.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) CheckS3Access(ctx context.Context, req *agentpb.CheckS3AccessReq) (*agentpb.CheckS3AccessResp, error) {
	url := fmt.Sprintf("%s/api/v1/s3/check_access", c.AgentUrl)
	respBytes, err := c.getResponseProtoJson(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp agentpb.CheckS3AccessResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) CheckAliOSSAccess(ctx context.Context, req *agentpb.CheckAliOSSAccessReq) (*agentpb.CheckAliOSSAccessResp, error) {
	url := fmt.Sprintf("%s/api/v1/alioss/check_access", c.AgentUrl)
	respBytes, err := c.getResponseProtoJson(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp agentpb.CheckAliOSSAccessResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}

func (c *NGRegionalClient) CheckAzblobAccess(ctx context.Context, req *agentpb.CheckAzblobAccessReq) (*agentpb.CheckAzblobAccessResp, error) {
	url := fmt.Sprintf("%s/api/v1/azblob/check_access", c.AgentUrl)
	respBytes, err := c.getResponseProtoJson(ctx, http.MethodPost, url, req)
	if err != nil {
		return nil, err
	}
	var resp agentpb.CheckAzblobAccessResp
	if err := protojson.Unmarshal(respBytes, &resp); err != nil {
		return nil, errors.Trace(err)
	}
	return &resp, nil
}