package pimport

import (
	"bytes"
	"compress/gzip"
	"context"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/pingcap/errcode"
	"github.com/pingcap/errors"
	"github.com/tidbcloud/core-svc/idl/pbgen/proto/core"
	"github.com/tidbcloud/dataflow-service/internal/db/entity"
	"github.com/tidbcloud/dataflow-service/internal/openapi/service"
	"github.com/tidbcloud/dataflow-service/internal/utils"
	"github.com/tidbcloud/idl/pbgen/proto/openapi/dp"
	"github.com/tidbcloud/serverless-service/idl/pbgen/proto/console"
	tidb_mgmt_servicev1 "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
)

func (s *Service) ListServerlessImports(ctx context.Context, cluster *tidb_mgmt_servicev1.Cluster, pager utils.Pager, req *dp.ListImportTasksReq, userContext *service.UserContext) (*dp.ListImportTasksResp, error) {
	err := checkPermission(userContext, cluster.OrganizationId, cluster.ProjectId, req.GetProjectId())
	if err != nil {
		return nil, err
	}

	resp, err := s.serverlessSvcGrpcClient.ListImports(ctx, &console.ListImportsReq{
		OrgId:     strconv.FormatUint(cluster.OrganizationId, 10),
		ProjectId: strconv.FormatUint(cluster.ProjectId, 10),
		ClusterId: strconv.FormatUint(cluster.ClusterId, 10),
		PageSize:  pager.PageSize,
		PageToken: strconv.FormatInt(int64(pager.Page), 10),
	})
	if err != nil {
		return nil, err
	}

	imports := make([]*dp.ImportItem, 0, len(resp.Imports))
	for _, i := range resp.Imports {
		elem := convertToDpImport(i)
		imports = append(imports, elem)
	}

	return &dp.ListImportTasksResp{
		Items: imports,
		Total: resp.TotalSize,
	}, nil
}

func (s *Service) GetServerlessImport(ctx context.Context, cluster *tidb_mgmt_servicev1.Cluster, req *dp.GetImportTaskReq, userContext *service.UserContext) (*dp.ImportItem, error) {
	err := checkPermission(userContext, cluster.OrganizationId, cluster.ProjectId, req.GetProjectId())
	if err != nil {
		return nil, err
	}

	i, err := s.serverlessSvcGrpcClient.GetImport(ctx, &console.ImportIds{
		OrgId:     strconv.FormatUint(cluster.OrganizationId, 10),
		ProjectId: strconv.FormatUint(cluster.ProjectId, 10),
		ClusterId: strconv.FormatUint(cluster.ClusterId, 10),
		Id:        strconv.FormatUint(req.GetImportId(), 10),
	})
	if err != nil {
		return nil, err
	}

	return convertToDpImport(i), nil
}

func (s *Service) GetDataIngestionProfile(ctx context.Context, cluster *tidb_mgmt_servicev1.Cluster, req *dp.GetImportTaskRoleInfoReq, userContext *service.UserContext) (*dp.ImportTaskRoleInfo, error) {
	err := checkPermission(userContext, cluster.OrganizationId, cluster.ProjectId, req.GetProjectId())
	if err != nil {
		return nil, err
	}

	profile, err := s.serverlessSvcGrpcClient.GetDataIngestionProfile(ctx, &console.GetProfileReq{
		OrgId:     strconv.FormatUint(cluster.OrganizationId, 10),
		ProjectId: strconv.FormatUint(cluster.ProjectId, 10),
		ClusterId: strconv.FormatUint(cluster.ClusterId, 10),
	})
	if err != nil {
		return nil, err
	}

	return &dp.ImportTaskRoleInfo{
		RoleInfo: &dp.ImportTaskRoleInfo_AwsImportRole{
			AwsImportRole: &dp.AwsImportTaskRoleInfo{
				AccountId:  profile.GetS3Profile().GetAwsAccountId(),
				ExternalId: profile.GetS3Profile().GetAwsExternalId(),
			},
		},
	}, nil
}

func (s *Service) UpdateServerlessImport(ctx context.Context, cluster *tidb_mgmt_servicev1.Cluster, req *dp.UpdateImportTaskReq, userContext *service.UserContext) error {
	err := checkPermission(userContext, cluster.OrganizationId, cluster.ProjectId, req.GetProjectId())
	if err != nil {
		return err
	}

	if req.GetAction() == dp.UpdateImportTaskReq_CANCEL {
		err = s.serverlessSvcGrpcClient.CancelImport(ctx, &console.ImportIds{
			OrgId:     strconv.FormatUint(cluster.OrganizationId, 10),
			ProjectId: strconv.FormatUint(cluster.ProjectId, 10),
			ClusterId: strconv.FormatUint(cluster.ClusterId, 10),
			Id:        strconv.FormatUint(req.GetImportId(), 10),
		})
		if err != nil {
			return err
		}
	}

	return nil
}

func (s *Service) CreateServerlessImport(ctx context.Context, cluster *tidb_mgmt_servicev1.Cluster, req *dp.CreateImportTaskReq, userContext *service.UserContext) (*dp.CreateImportTaskResp, error) {
	err := checkPermission(userContext, cluster.OrganizationId, cluster.ProjectId, req.GetProjectId())
	if err != nil {
		return nil, err
	}

	importReq := &console.CreateImportReq{
		OrgId:     strconv.FormatUint(cluster.OrganizationId, 10),
		ProjectId: strconv.FormatUint(cluster.ProjectId, 10),
		ClusterId: strconv.FormatUint(cluster.ClusterId, 10),
		// apikey-<publicKey> is used as the creator name
		Creator: fmt.Sprintf("apikey-%s", userContext.User().LastName),
	}
	if req.GetSpec() == nil {
		return nil, errcode.NewInvalidInputErr(errors.New("spec is required"))
	}
	if req.GetSpec().GetSource() == nil {
		return nil, errcode.NewInvalidInputErr(errors.New("source is required"))
	}
	uri := req.GetSpec().GetSource().GetUri()
	if req.GetSpec().GetSource().GetType() == dp.ImportSource_S3 {
		if uri == "" {
			return nil, errcode.NewInvalidInputErr(errors.New("source uri is required"))
		}

		importReq.Source = &console.ImportSource{
			Type: console.ImportSourceType_S3,
			Source: &console.ImportSource_S3{
				S3: &console.S3Source{
					Uri: uri,
				},
			},
		}
		if req.GetSpec().GetSource().GetAwsAssumeRoleAccess() != nil {
			importReq.Source.GetS3().AuthType = console.ImportS3AuthType_ROLE_ARN
			role := req.GetSpec().GetSource().GetAwsAssumeRoleAccess().GetAssumeRole()
			if role == "" {
				return nil, errcode.NewInvalidInputErr(errors.New("role arn is required"))
			}
			importReq.Source.GetS3().Auth = &console.S3Source_RoleArn{
				RoleArn: role,
			}
		} else if req.GetSpec().GetSource().GetAwsKeyAccess() != nil {
			importReq.Source.GetS3().AuthType = console.ImportS3AuthType_ACCESS_KEY
			key := req.GetSpec().GetSource().GetAwsKeyAccess()
			if key.GetAccessKeyId() == "" {
				return nil, errcode.NewInvalidInputErr(errors.New("access key id is required"))
			}
			if key.GetSecretAccessKey() == "" {
				return nil, errcode.NewInvalidInputErr(errors.New("secret access key is required"))
			}
			importReq.Source.GetS3().Auth = &console.S3Source_AccessKey_{
				AccessKey: &console.S3Source_AccessKey{
					Id:     key.GetAccessKeyId(),
					Secret: key.GetSecretAccessKey(),
				},
			}
		} else {
			return nil, errcode.NewInvalidInputErr(errors.New("auth is required"))
		}

		if req.GetSpec().GetTarget() != nil {
			tables := make([]*console.ImportTargetTableInfo, 0, len(req.GetSpec().GetTarget().GetTables()))
			for _, t := range req.GetSpec().GetTarget().GetTables() {
				tables = append(tables, &console.ImportTargetTableInfo{
					TargetTable: &console.Table{
						Database: t.GetDatabaseName(),
						Table:    t.GetTableName(),
					},
					CustomFile: t.GetFileNamePattern(),
				})
			}
		}
	} else if req.GetSpec().GetSource().GetType() == dp.ImportSource_LOCAL_FILE {
		importReq.Source = &console.ImportSource{
			Type: console.ImportSourceType_LOCAL,
			Source: &console.ImportSource_Local{
				Local: &console.LocalSource{
					Columns:   nil,
					PkColumns: nil,
				},
			},
		}
		if uri == "" {
			return nil, errcode.NewInvalidInputErr(errors.New("upload stub id is required"))
		}
		u, err := url.Parse(uri)
		if err != nil {
			return nil, errcode.NewInvalidInputErr(errors.New("invalid uri format"))
		}
		importReq.Source.GetLocal().UploadId = u.Host
		if req.GetSpec().GetTarget() == nil {
			return nil, errcode.NewInvalidInputErr(errors.New("target is required for local file import"))
		}
		if len(req.GetSpec().GetTarget().GetTables()) != 1 {
			return nil, errcode.NewInvalidInputErr(errors.New("only one target table is supported when the source type is `LOCAL_FILE`"))
		}
		importReq.Source.GetLocal().TargetDatabase = req.GetSpec().GetTarget().GetTables()[0].GetDatabaseName()
		importReq.Source.GetLocal().TargetTable = req.GetSpec().GetTarget().GetTables()[0].GetTableName()
	} else {
		return nil, errcode.NewInvalidInputErr(errors.New("unsupported source type for serverless import: " + req.GetSpec().GetSource().GetType().String()))
	}

	if req.GetSpec().GetSource().GetFormat() == nil {
		return nil, errcode.NewInvalidInputErr(errors.New("format is required"))
	}
	importReq.ImportOptions = &console.ImportOptions{
		FileType:    convertToRpcFileType(req.GetSpec().GetSource().GetFormat().GetType()),
		DisplayName: req.GetName(),
	}
	if req.GetSpec().GetSource().GetFormat().GetFormatConfig() != nil {
		csv := req.GetSpec().GetSource().GetFormat().GetCsvConfig()
		if csv != nil {
			importReq.ImportOptions.Format = &console.ImportOptions_CsvFormat{
				CsvFormat: &console.CSVFormat{
					Separator:       csv.GetDelimiter(),
					Delimiter:       aws.String(csv.GetQuote()),
					Header:          aws.Bool(csv.GetHasHeaderRow()),
					BackslashEscape: aws.Bool(csv.GetBackslashEscape()),
				},
			}
		}
	}
	i, err := s.serverlessSvcGrpcClient.CreateImport(ctx, importReq)
	if err != nil {
		return nil, err
	}

	return &dp.CreateImportTaskResp{
		Id: i.CompatibleId,
	}, nil
}

func convertToRpcFileType(formatType dp.ImportSourceFormat_ImportSourceFormatType) console.ImportFileType_Enum {
	switch formatType {
	case dp.ImportSourceFormat_CSV:
		return console.ImportFileType_CSV
	case dp.ImportSourceFormat_PARQUET:
		return console.ImportFileType_PARQUET
	case dp.ImportSourceFormat_SQL:
		return console.ImportFileType_SQL
	case dp.ImportSourceFormat_AURORA_SNAPSHOT:
		return console.ImportFileType_AURORA_SNAPSHOT
	default:
		return console.ImportFileType_UNSPECIFIED
	}
}

func (s *Service) UploadLoadFileForServerless(ctx context.Context, cluster *core.Cluster, req *dp.UploadLocalFileReq, userContext *service.UserContext) (*dp.UploadLocalFileResp, error) {
	err := checkPermission(userContext, cluster.TenantId, cluster.ProjectId, req.GetProjectId())
	if err != nil {
		return nil, err
	}

	upload, err := s.serverlessSvcGrpcClient.Upload(ctx, &console.UploadRequest{
		OrgId:     strconv.FormatUint(cluster.TenantId, 10),
		ProjectId: strconv.FormatUint(cluster.ProjectId, 10),
		ClusterId: strconv.FormatUint(cluster.Id, 10),
		FileName:  req.LocalFileName,
	})
	if err != nil {
		return nil, err
	}

	payload := req.GetPayload()
	if payload == nil {
		return nil, errcode.NewInvalidInputErr(errors.New("payload is not specified"))
	}

	err = uploadFile(ctx, upload.UploadUrl, payload)
	if err != nil {
		return nil, err
	}

	return &dp.UploadLocalFileResp{
		UploadStubId: upload.CompatibleId,
	}, nil
}

func (s *Service) PreviewServerlessImportData(ctx context.Context, cluster *tidb_mgmt_servicev1.Cluster, req *dp.PreviewImportDataReq, userContext *service.UserContext) (*dp.PreviewImportDataResp, error) {
	err := checkPermission(userContext, cluster.OrganizationId, cluster.ProjectId, req.GetProjectId())
	if err != nil {
		return nil, err
	}

	if req.GetSpec() == nil {
		return nil, errcode.NewInvalidInputErr(errors.New("spec is required"))
	}
	if req.GetSpec().GetSource() == nil {
		return nil, errcode.NewInvalidInputErr(errors.New("source is required"))
	}
	if req.GetSpec().GetSource().GetType() != dp.ImportSource_LOCAL_FILE {
		return nil, errcode.NewInvalidInputErr(errors.New("only local file import is supported for data preview"))
	}
	if req.GetSpec().GetSource().GetFormat() == nil {
		return nil, errcode.NewInvalidInputErr(errors.New("format is required"))
	}

	localReq := &console.ImportsPreviewLocalReq{
		OrgId:     strconv.FormatUint(cluster.OrganizationId, 10),
		ProjectId: strconv.FormatUint(cluster.ProjectId, 10),
		ClusterId: strconv.FormatUint(cluster.ClusterId, 10),
		UploadId:  req.GetSpec().GetSource().GetUri(),
		FileType:  convertToRpcFileType(req.GetSpec().GetSource().GetFormat().GetType()),
		CsvFormat: nil,
	}

	uri := req.GetSpec().GetSource().GetUri()
	if uri == "" {
		return nil, errcode.NewInvalidInputErr(errors.New("upload stub id is required"))
	}
	u, err := url.Parse(uri)
	if err != nil {
		return nil, errcode.NewInvalidInputErr(errors.New("invalid uri format"))
	}
	localReq.UploadId = u.Host
	if req.GetSpec().GetSource().GetFormat() != nil {
		csv := req.GetSpec().GetSource().GetFormat().GetCsvConfig()
		if csv != nil {
			localReq.CsvFormat = &console.CSVFormat{
				Separator:       csv.GetDelimiter(),
				Delimiter:       aws.String(csv.GetQuote()),
				BackslashEscape: aws.Bool(csv.GetBackslashEscape()),
				Header:          aws.Bool(csv.GetHasHeaderRow()),
			}
		}
	}
	preview, err := s.serverlessSvcGrpcClient.PreviewLocal(ctx, localReq)
	if err != nil {
		return nil, err
	}

	return &dp.PreviewImportDataResp{
		TablePreviews: convertToDpPreview(preview),
	}, nil
}

func convertToDpPreview(preview *console.ImportsPreviewLocalResp) []*dp.TablePreview {
	elem := &dp.TablePreview{
		DatabaseName:  preview.GetTable().GetDatabaseName(),
		TableName:     preview.GetTable().GetTableName(),
		SchemaPreview: &dp.TableSchema{},
		DataPreview:   &dp.TableData{},
	}
	columns := make([]*dp.ColumnDefinition, 0, len(preview.GetTable().GetRowDataTypes()))
	columnsValues := make([]*dp.TableDataRow, 0, len(preview.GetTable().GetRowDataTypes()))
	for i := range preview.GetTable().GetRowDataTypes() {
		columns = append(columns, &dp.ColumnDefinition{
			ColumnName: preview.GetTable().GetRowHeaders()[i],
			ColumnType: preview.GetTable().GetRowDataTypes()[i],
		})
		columnsValues = append(columnsValues, &dp.TableDataRow{
			Columns: preview.GetTable().GetRowValues()[i].GetColumnValues(),
		})
	}
	elem.SchemaPreview.ColumnDefinitions = columns
	elem.DataPreview.ColumnNames = preview.GetTable().GetRowHeaders()
	elem.DataPreview.Rows = columnsValues

	return []*dp.TablePreview{elem}
}

func uploadFile(ctx context.Context, uploadUrl string, payload *dp.LocalFilePayload) error {
	rawReader := bytes.NewReader(payload.GetContent())
	actualContentReader, err := gzip.NewReader(rawReader)
	if err != nil {
		return err
	}
	//nolint: errcheck
	defer actualContentReader.Close()

	req, err := http.NewRequestWithContext(
		ctx,
		http.MethodPut,
		uploadUrl,
		io.LimitReader(actualContentReader, int64(payload.GetTotalSizeBytes())),
	)
	if err != nil {
		return errors.Annotate(err, "construct HTTP request error")
	}
	req.ContentLength = int64(payload.GetTotalSizeBytes())
	client := http.DefaultClient
	resp, err := client.Do(req)
	if err != nil {
		return errors.Annotate(err, "send request to the target S3 upload URL failed")
	}
	//nolint:errcheck
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		body := ""
		var b strings.Builder
		if _, err := io.Copy(&b, resp.Body); err == nil {
			body = b.String()
		}
		return fmt.Errorf("PUT data to target has error: unexpected status code: %d.  Body: %s", resp.StatusCode, body)
	}
	return nil
}

func checkPermission(userContext *service.UserContext, clusterTenantID uint64, clusterProjectID uint64,
	reqProjectID uint64) error {
	bizErr := userContext.CheckOrgIDAndProjectID(entity.ID(clusterTenantID), entity.ID(clusterProjectID))
	if bizErr != nil {
		return bizErr
	}
	if clusterProjectID != reqProjectID {
		return errcode.NewForbiddenErr(errors.New("you do not have permission to access this resource"))
	}
	return nil
}

func convertToDpImport(i *console.Import) *dp.ImportItem {
	elem := &dp.ImportItem{
		Metadata: &dp.ImportMetadata{
			Id:   i.CompatibleId,
			Name: i.DisplayName,
		},
		Spec: &dp.ImportSpec{
			Source: &dp.ImportSource{
				Type: convertToDpType(i.GetCreationDetails().GetSource().GetType()),
				Format: &dp.ImportSourceFormat{
					Type: convertToDpFormat(i.GetCreationDetails().GetImportOptions().GetFileType()),
				},
			},
			Target: &dp.ImportTarget{},
		},
		Status: &dp.ImportStatus{
			Phase:        convertToPhase(i.State),
			ErrorMessage: i.Message,
			Progress: &dp.ImportProgress{
				ImportProgress:     float64(i.CompletePercent),
				ValidationProgress: float64(i.PostImportCompletedPercent),
			},
			SourceTotalSizeBytes: i.TotalSize,
		},
	}
	if i.CompleteTime != nil {
		elem.Status.EndTimestamp = i.CompleteTime.AsTime().Unix()
	}
	if i.StartTime != nil {
		elem.Status.StartTimestamp = i.StartTime.AsTime().Unix()
	}
	if i.CreateTime != nil {
		elem.Metadata.CreateTimestamp = i.CreateTime.AsTime().Unix()
	}
	format := i.GetCreationDetails().GetImportOptions().GetCsvFormat()

	if i.GetCreationDetails().GetImportOptions().GetFileType() == console.ImportFileType_CSV {
		if format != nil {
			elem.Spec.Source.Format.FormatConfig = &dp.ImportSourceFormat_CsvConfig{
				CsvConfig: &dp.ImportSourceCSVConfig{
					Delimiter:       format.GetSeparator(),
					Quote:           format.GetDelimiter(),
					BackslashEscape: format.GetBackslashEscape(),
					HasHeaderRow:    format.GetHeader(),
				},
			}
		} else {
			elem.Spec.Source.Format.FormatConfig = &dp.ImportSourceFormat_CsvConfig{
				CsvConfig: &dp.ImportSourceCSVConfig{
					Delimiter:       ",",
					Quote:           "\"",
					BackslashEscape: true,
					HasHeaderRow:    true,
				},
			}
		}
	}

	if i.GetCreationDetails().GetSource().GetType() == console.ImportSourceType_S3 {
		s3 := i.GetCreationDetails().GetSource().GetS3()
		infos := s3.GetTargetTableInfos()
		tables := make([]*dp.ImportTargetTable, 0, len(infos))
		for _, info := range infos {
			tables = append(tables, &dp.ImportTargetTable{
				DatabaseName:    info.TargetTable.Database,
				TableName:       info.TargetTable.Table,
				FileNamePattern: info.CustomFile,
			})
		}
		elem.Spec.Target.Tables = tables
		elem.Spec.Source.Uri = s3.GetUri()
		if s3.GetAuthType() == console.ImportS3AuthType_ACCESS_KEY {
			elem.Spec.Source.Access = &dp.ImportSource_AwsKeyAccess{
				AwsKeyAccess: &dp.AwsKeyAccess{
					AccessKeyId:     "REDACTED",
					SecretAccessKey: "REDACTED",
				},
			}
		} else if s3.GetAuthType() == console.ImportS3AuthType_ROLE_ARN {
			elem.Spec.Source.Access = &dp.ImportSource_AwsAssumeRoleAccess{
				AwsAssumeRoleAccess: &dp.AwsAssumeRoleAccess{
					AssumeRole: s3.GetRoleArn(),
				},
			}
		}
	} else if i.GetCreationDetails().GetSource().GetType() == console.ImportSourceType_LOCAL {
		elem.Spec.Target.Tables = []*dp.ImportTargetTable{
			{
				DatabaseName: i.GetCreationDetails().GetSource().GetLocal().GetTargetDatabase(),
				TableName:    i.GetCreationDetails().GetSource().GetLocal().GetTargetTable(),
			},
		}
	}
	return elem
}

func convertToDpFormat(fileType console.ImportFileType_Enum) dp.ImportSourceFormat_ImportSourceFormatType {
	switch fileType {
	case console.ImportFileType_CSV:
		return dp.ImportSourceFormat_CSV
	case console.ImportFileType_PARQUET:
		return dp.ImportSourceFormat_PARQUET
	case console.ImportFileType_SQL:
		return dp.ImportSourceFormat_SQL
	case console.ImportFileType_AURORA_SNAPSHOT:
		return dp.ImportSourceFormat_AURORA_SNAPSHOT
	default:
		return dp.ImportSourceFormat_UNSPECIFIED
	}
}

func convertToDpType(getType console.ImportSourceType_Enum) dp.ImportSource_ImportSourceType {
	switch getType {
	case console.ImportSourceType_S3:
		return dp.ImportSource_S3
	case console.ImportSourceType_LOCAL:
		return dp.ImportSource_LOCAL_FILE
	case console.ImportSourceType_GCS:
		return dp.ImportSource_GCS
	default:
		return dp.ImportSource_UNSPECIFIED
	}
}

func convertToPhase(state console.ImportState_Enum) dp.ImportStatus_ImportTaskPhase {
	switch state {
	case console.ImportState_PREPARING:
		return dp.ImportStatus_PREPARING
	case console.ImportState_IMPORTING:
		return dp.ImportStatus_IMPORTING
	case console.ImportState_COMPLETED:
		return dp.ImportStatus_COMPLETED
	case console.ImportState_FAILED:
		return dp.ImportStatus_FAILED
	case console.ImportState_CANCELING:
		return dp.ImportStatus_CANCELING
	case console.ImportState_CANCELED:
		return dp.ImportStatus_CANCELED
	default:
		return dp.ImportStatus_UNSPECIFIED
	}
}