package schsdk

import (
	"gitlink.org.cn/cloudream/common/pkgs/types"
	cdssdk "gitlink.org.cn/cloudream/common/sdks/storage"
	"gitlink.org.cn/cloudream/common/utils/serder"
)

const (
	JobTypeNormal         = "Normal"
	JobTypePCM            = "PCM"
	JobTypeResource       = "Resource"
	JobTypeInstance       = "Instance"
	JobTypeFinetuning     = "Finetuning"
	JobTypeDataPreprocess = "DataPreprocess"
	JobTypeDataReturn     = "DataReturn"

	FileInfoTypePackage   = "Package"
	FileInfoTypeLocalFile = "LocalFile"
	FileInfoTypeResource  = "Resource"
	FileInfoTypeImage     = "Image"

	FILE   = "file"
	FOLDER = "folder"

	MemoryUtilization = "MemoryUtilization"
	GPUUtilization    = "GPUUtilization"
	CPUUtilization    = "CPUUtilization"

	MethodPost  = "POST"
	MethodGet   = "GET"
	CodeSuccess = 200
)

type JobID string

type JobSetID string

type DataID int64

type ImageID int64

// 计算中心ID
type CCID int64

type ModelID string
type ModelName string
type ECSInstanceID string
type NodeID int64
type Address string

type ClusterID string

type JobSetInfo struct {
	Jobs []JobInfo `json:"jobs"`
}

type JobInfo interface {
	GetLocalJobID() string
	GetTargetLocalJobIDs() []string
	SetTargetLocalJob(info TargetJobInfo)
	GetTargetInputParams(targetID string) map[string]string
}

var JobInfoTypeUnion = types.NewTypeUnion[JobInfo](
	(*NormalJobInfo)(nil),
	(*DataReturnJobInfo)(nil),
	(*MultiInstanceJobInfo)(nil),
	(*InstanceJobInfo)(nil),
	(*UpdateMultiInstanceJobInfo)(nil),
	(*FinetuningJobInfo)(nil),
	(*DataPreprocessJobInfo)(nil),
	(*AIJobInfo)(nil),
	(*HPCJobInfo)(nil),
	(*BindingJobInfo)(nil),
	(*PCMInferenceJobInfo)(nil),
	(*CompleteJobInfo)(nil),
	(*StartJobInfo)(nil),
	(*NotifyJobInfo)(nil),
	(*StopInferenceJobInfo)(nil),
)
var _ = serder.UseTypeUnionInternallyTagged(&JobInfoTypeUnion, "type")

type JobInfoBase struct {
	LocalJobID string          `json:"localJobID"`
	TargetJob  []TargetJobInfo `json:"targetJob"`
}

type TargetJobInfo struct {
	TargetJobID string            `json:"targetJobID"`
	InputParams map[string]string `json:"inputParams"`
}

func (i *JobInfoBase) GetLocalJobID() string {
	return i.LocalJobID
}

func (i *JobInfoBase) GetTargetInputParams(targetID string) map[string]string {
	for _, v := range i.TargetJob {
		if v.TargetJobID == targetID {
			return v.InputParams
		}
	}
	return nil
}

func (i *JobInfoBase) GetTargetLocalJobIDs() []string {
	var IDs []string
	for _, v := range i.TargetJob {
		IDs = append(IDs, v.TargetJobID)
	}
	return IDs
}

func (i *JobInfoBase) SetTargetLocalJob(info TargetJobInfo) {
	for _, target := range i.TargetJob {
		// 已经存在，则不用再添加
		if target.TargetJobID == info.TargetJobID {
			return
		}
	}
	i.TargetJob = append(i.TargetJob, info)
}

type NormalJobInfo struct {
	serder.Metadata `union:"Normal"`
	JobInfoBase
	Type         string           `json:"type"`
	Files        JobFilesInfo     `json:"files"`
	Runtime      JobRuntimeInfo   `json:"runtime"`
	Resources    JobResourcesInfo `json:"resources"`
	Services     JobServicesInfo  `json:"services"`
	ModelJobInfo ModelJobInfo     `json:"modelJobInfo"`
}

type PCMInferenceJobInfo struct {
	serder.Metadata `union:"PCM_Inference"`
	JobInfoBase
	Type           string         `json:"type"`
	Name           string         `json:"name"`
	Description    string         `json:"description"`
	Files          JobFilesInfo   `json:"files"`
	JobResources   JobResources   `json:"jobResources"`
	BindingID      DataID         `json:"bindingID"`
	ResourceChoice ResourceChoice `json:"resourceChoice"`
}

type StopInferenceJobInfo struct {
	serder.Metadata `union:"StopInference"`
	JobInfoBase
	Type string `json:"type"`
	Url  string `json:"url"`
}

type AIJobInfo struct {
	serder.Metadata `union:"AI"`
	JobInfoBase
	Type           string         `json:"type"`
	Name           string         `json:"name"`
	Description    string         `json:"description"`
	Files          JobFilesInfo   `json:"files"`
	JobResources   JobResources   `json:"jobResources"`
	ResourceChoice ResourceChoice `json:"resourceChoice"`
}

type CompleteJobInfo struct {
	serder.Metadata `union:"Complete"`
	JobInfoBase
	Type string `json:"type"`
}

type StartJobInfo struct {
	serder.Metadata `union:"Start"`
	JobInfoBase
	Type string `json:"type"`
}

type NotifyJobInfo struct {
	serder.Metadata `union:"Notify"`
	JobInfoBase
	Type        string            `json:"type"`
	RequestType string            `json:"requestType"`
	Url         string            `json:"url"`
	Body        any               `json:"body"`
	Headers     map[string]string `json:"headers"`
}

type ResourceChoice struct {
	Type           string          `json:"type"`
	ResourceScopes []ResourceScope `json:"resourceScopes"`
}

type ResourceScope struct {
	Name string  `json:"name"`
	Min  float64 `json:"min"`
	Max  float64 `json:"max"`
}

type BindingJobInfo struct {
	serder.Metadata `union:"Binding"`
	JobInfoBase
	Type string      `json:"type"`
	Info DataBinding `json:"info"`
	Name string      `json:"name"` // 临时使用
}

type DataBinding interface {
	Noop()
}

var DataBindingTypeUnion = types.NewTypeUnion[DataBinding](
	(*ModelBinding)(nil),
)

var _ = serder.UseTypeUnionInternallyTagged(&DataBindingTypeUnion, "type")

type DataBindingBase struct{}

func (d *DataBindingBase) Noop() {}

type ModelBinding struct {
	serder.Metadata `union:"model"`
	DataBindingBase
	Type           string      `json:"type"`
	Name           string      `json:"name"`
	Description    string      `json:"description"`
	ClusterIDs     []ClusterID `json:"clusterIDs"`
	Category       string      `json:"category"`
	ModelType      string      `json:"modelType"`
	Env            string      `json:"env"`
	Version        string      `json:"version"`
	RepositoryName string      `json:"repositoryName"`
}

type HPCJobInfo struct {
	serder.Metadata `union:"HPC"`
	JobInfoBase
	Type          string       `json:"type"`
	Name          string       `json:"name"`
	Description   string       `json:"description"`
	ClusterID     ClusterID    `json:"clusterID"`
	Backend       string       `json:"backend"`
	App           string       `json:"app"`
	OperateType   string       `json:"operateType"`
	ScriptContent string       `json:"scriptContent"`
	Parameters    HPCParameter `json:"parameters"`
}

type HPCParameter struct {
	JobName         string           `json:"jobName"`
	JobDir          string           `json:"jobDir"`
	Partition       string           `json:"partition"`
	Ntasks          string           `json:"ntasks"`
	Nodes           string           `json:"nodes"`
	BamFile         string           `json:"bamFile"`
	HashType        string           `json:"hashType"`
	AttackMode      string           `json:"attackMode"`
	HashInput       string           `json:"hashInput"`
	Mask            string           `json:"mask"`
	Dictionary      string           `json:"dictionary"`
	Dictionary2     string           `json:"dictionary2"`
	HPCBindingFiles []HPCBindingFile `json:"hpcBindingFiles"`
}

type HPCBindingFile struct {
	ParamName string  `json:"paramName"`
	Resource  HPCFile `json:"resource"`
}

type HPCFile interface {
	Noop()
}

var HPCFileTypeUnion = types.NewTypeUnion[HPCFile](
	(*HPCObject)(nil),
	(*HPCPath)(nil),
)

var _ = serder.UseTypeUnionInternallyTagged(&HPCFileTypeUnion, "type")

type HPCFileBase struct{}

func (d *HPCFileBase) Noop() {}

type HPCObject struct {
	serder.Metadata `union:"object"`
	HPCFileBase
	Type     string          `json:"type"`
	ObjectID cdssdk.ObjectID `json:"objectID"`
}

type HPCPath struct {
	serder.Metadata `union:"path"`
	HPCFileBase
	Type      string           `json:"type"`
	PackageID cdssdk.PackageID `json:"packageID"`
	Path      string           `json:"path"`
}

type JobResources struct {
	//任务分配策略：负载均衡、积分优先、随机分配等，dataLocality, leastLoadFirst
	ScheduleStrategy string        `json:"scheduleStrategy"`
	Clusters         []ClusterInfo `json:"clusters"`
}

type ClusterInfo struct {
	ClusterID ClusterID     `json:"clusterID"`
	Resources []JobResource `json:"resources"`
	//Files     JobFilesInfo      `json:"files"`
	Code    JobFileInfo       `json:"code"`
	Runtime PCMJobRuntimeInfo `json:"runtime"`
}

type PCMJobRuntimeInfo struct {
	Command string                 `json:"command"`
	Envs    map[string]interface{} `json:"envs"`
	Params  map[string]interface{} `json:"params"`
}

//type Resource struct {
//	Resource []JobResource `json:"resource"`
//}

type JobResource interface {
	Noop()
}

var JobResourceTypeUnion = types.NewTypeUnion[JobResource](
	(*CPU)(nil),
	(*GPU)(nil),
	(*NPU)(nil),
	(*MLU)(nil),
	(*DCU)(nil),
	(*MEMORY)(nil),
	(*PRICE)(nil),
	(*STORAGE)(nil),
)

var _ = serder.UseTypeUnionInternallyTagged(&JobResourceTypeUnion, "type")

type JobResourceBase struct{}

func (d *JobResourceBase) Noop() {}

type CPU struct {
	serder.Metadata `union:"CPU"`
	JobResourceBase
	Type   string `json:"type"`
	Name   string `json:"name"`
	Number int64  `json:"number"`
}

type STORAGE struct {
	serder.Metadata `union:"STORAGE"`
	JobResourceBase
	Type   string `json:"type"`
	Name   string `json:"name"`
	Number int64  `json:"number"`
}

type GPU struct {
	serder.Metadata `union:"GPU"`
	JobResourceBase
	Type   string `json:"type"`
	Name   string `json:"name"`
	Number int64  `json:"number"`
}

type NPU struct {
	serder.Metadata `union:"NPU"`
	JobResourceBase
	Type   string `json:"type"`
	Name   string `json:"name"`
	Number int64  `json:"number"`
}

type MEMORY struct {
	serder.Metadata `union:"MEMORY"`
	JobResourceBase
	Type   string `json:"type"`
	Name   string `json:"name"`
	Number int64  `json:"number"`
}

type DCU struct {
	serder.Metadata `union:"DCU"`
	JobResourceBase
	Type   string `json:"type"`
	Name   string `json:"name"`
	Number int64  `json:"number"`
}

type MLU struct {
	serder.Metadata `union:"MLU"`
	JobResourceBase
	Type   string `json:"type"`
	Name   string `json:"name"`
	Number int64  `json:"number"`
}

type PRICE struct {
	serder.Metadata `union:"PRICE"`
	JobResourceBase
	Type   string `json:"type"`
	Name   string `json:"name"`
	Number int64  `json:"number"`
}

// FinetuningJobInfo 模型微调
type FinetuningJobInfo struct {
	serder.Metadata `union:"Finetuning"`
	JobInfoBase
	Type         string           `json:"type"`
	Files        JobFilesInfo     `json:"files"`
	Runtime      JobRuntimeInfo   `json:"runtime"`
	Resources    JobResourcesInfo `json:"resources"`
	Services     JobServicesInfo  `json:"services"`
	ModelJobInfo ModelJobInfo     `json:"modelJobInfo"`
}

// DataPreprocessJobInfo 数据预处理
type DataPreprocessJobInfo struct {
	serder.Metadata `union:"DataPreprocess"`
	JobInfoBase
	Type      string           `json:"type"`
	Files     JobFilesInfo     `json:"files"`
	Runtime   JobRuntimeInfo   `json:"runtime"`
	Resources JobResourcesInfo `json:"resources"`
	Services  JobServicesInfo  `json:"services"`
}

type DataReturnJobInfo struct {
	serder.Metadata `union:"DataReturn"`
	JobInfoBase
	Type             string               `json:"type"`
	BucketID         cdssdk.BucketID      `json:"bucketID"`
	TargetLocalJobID string               `json:"targetLocalJobID"`
	ReportMessage    TrainJobStatusReport `json:"report"`
}

// MultiInstanceJobInfo 多实例(推理任务)
type MultiInstanceJobInfo struct {
	serder.Metadata `union:"MultiInstance"`
	JobInfoBase
	Type         string           `json:"type"`
	Files        JobFilesInfo     `json:"files"`
	Runtime      JobRuntimeInfo   `json:"runtime"`
	Resources    JobResourcesInfo `json:"resources"`
	ModelJobInfo ModelJobInfo     `json:"modelJobInfo"`
}

// UpdateMultiInstanceJobInfo 更新模型
type UpdateMultiInstanceJobInfo struct {
	serder.Metadata `union:"UpdateModel"`
	JobInfoBase
	Type                  string         `json:"type"`
	Files                 JobFilesInfo   `json:"files"`
	Runtime               JobRuntimeInfo `json:"runtime"`
	MultiInstanceJobSetID JobSetID       `json:"multiInstanceJobSetID"`
	UpdateType            string         `json:"updateType"`
	SubJobs               []JobID        `json:"subJobs"`
	Operate               string         `json:"operate"`
}

type ModelJobInfo struct {
	Type            string    `json:"type"`
	ModelID         ModelID   `json:"modelID"`
	CustomModelName ModelName `json:"customModelName"`
	Command         string    `json:"command"`
}

// InstanceJobInfo 单实例(推理任务)
type InstanceJobInfo struct {
	serder.Metadata `union:"Instance"`
	JobInfoBase
	Type         string           `json:"type"`
	LocalJobID   string           `json:"multiInstJobID"`
	Files        JobFilesInfo     `json:"files"`
	Runtime      JobRuntimeInfo   `json:"runtime"`
	Resources    JobResourcesInfo `json:"resources"`
	ModelJobInfo ModelJobInfo     `json:"modelJobInfo"`
}

type JobFilesInfo struct {
	Dataset JobFileInfo `json:"dataset"`
	Code    JobFileInfo `json:"code"`
	Image   JobFileInfo `json:"image"`
	Model   JobFileInfo `json:"model"`
}

type JobFileInfo interface {
	Noop()
}

var FileInfoTypeUnion = types.NewTypeUnion[JobFileInfo](
	(*PackageJobFileInfo)(nil),
	(*LocalJobFileInfo)(nil),
	(*DataReturnJobFileInfo)(nil),
	(*ImageJobFileInfo)(nil),
	(*BindingJobFileInfo)(nil),
)
var _ = serder.UseTypeUnionInternallyTagged(&FileInfoTypeUnion, "type")

type JobFileInfoBase struct{}

func (i *JobFileInfoBase) Noop() {}

type BindingJobFileInfo struct {
	serder.Metadata `union:"Binding"`
	JobFileInfoBase
	Type      string `json:"type"`
	BindingID int64  `json:"bindingID"`
	// 用于参数回显
	BindingName string `json:"bindingName"`
}

type PackageJobFileInfo struct {
	serder.Metadata `union:"Package"`
	JobFileInfoBase
	Type      string           `json:"type"`
	PackageID cdssdk.PackageID `json:"packageID"`
}

type LocalJobFileInfo struct {
	serder.Metadata `union:"LocalFile"`
	JobFileInfoBase
	Type      string `json:"type"`
	LocalPath string `json:"localPath"`
}

type DataReturnJobFileInfo struct {
	serder.Metadata `union:"DataReturn"`
	JobFileInfoBase
	Type                 string `json:"type"`
	DataReturnLocalJobID string `json:"dataReturnLocalJobID"`
}

type ImageJobFileInfo struct {
	serder.Metadata `union:"Image"`
	JobFileInfoBase
	Type    string  `json:"type"`
	ImageID ImageID `json:"imageID"`
	// 用于参数回显
	ImageName string `json:"imageName"`
}

type JobRuntimeInfo struct {
	Command string   `json:"command"`
	Envs    []KVPair `json:"envs"`
	Params  []KVPair `json:"params"`
}

type KVPair struct {
	Key   string `json:"key"`
	Value string `json:"value"`
}

// CPU、GPU、NPU、MLU单位为：核
// Storage、Memory单位为：字节
type JobResourcesInfo struct {
	CPU     float64 `json:"cpu"`
	GPU     float64 `json:"gpu"`
	NPU     float64 `json:"npu"`
	MLU     float64 `json:"mlu"`
	Storage int64   `json:"storage"`
	Memory  int64   `json:"memory"`
}

type JobSetFilesUploadScheme struct {
	LocalFileSchemes []LocalFileUploadScheme `json:"localFileUploadSchemes"`
}

type JobFilesUploadScheme struct {
	LocalFileSchemes []LocalFileUploadScheme `json:"localFileUploadSchemes"`
}

type LocalFileUploadScheme struct {
	LocalPath           string           `json:"localPath"`
	UploadToCDStorageID cdssdk.StorageID `json:"uploadToCDSStorageID"`
}

type JobServicesInfo struct {
	ServicePortInfos []ServicePortInfo `json:"servicePortInfos"`
}

type ServicePortInfo struct {
	Name string `json:"name"`
	Port int64  `json:"port"`
}

type JobSetServiceInfo struct {
	Name         string           `json:"name"`
	Port         int64            `json:"port"`
	CDSStorageID cdssdk.StorageID `json:"cdsStorageID"`
	LocalJobID   string           `json:"localJobID"`
}

type Bootstrap interface {
	GetBootstrapType() string
}

type DirectBootstrap struct {
	serder.Metadata `union:"Direct"`
	Type            string `json:"type"`
}

type NoEnvBootstrap struct {
	serder.Metadata `union:"NoEnv"`
	Type            string           `json:"type"`
	ScriptPackageID cdssdk.PackageID `json:"scriptPackageID"`
	ScriptFileName  string           `json:"scriptFileName"`
}

var BootstrapTypeUnion = types.NewTypeUnion[Bootstrap](
	(*DirectBootstrap)(nil),
	(*NoEnvBootstrap)(nil),
)

var _ = serder.UseTypeUnionInternallyTagged(&BootstrapTypeUnion, "type")

func (b *DirectBootstrap) GetBootstrapType() string {
	return b.Type
}

func (b *NoEnvBootstrap) GetBootstrapType() string {
	return b.Type
}

const (
	JobDataInEnv     = "SCH_DATA_IN"
	JobDataOutEnv    = "SCH_DATA_OUT"
	FinetuningOutEnv = "FINETUNING_OUT"
	AccessPath       = "ACCESS_PATH"
)

type Rclone struct {
	CDSRcloneID       string `json:"cds_rcloneID"`
	CDSRcloneConfigID string `json:"cds_rcloneConfigID"`
}

type InferencePlatform struct {
	PlatformName        string `json:"platformName"`
	ApiBaseUrl          string `json:"apiBaseUrl"`
	ApiKey              string `json:"apiKey"`
	ApiProxy            string `json:"apiProxy"`
	LlmModel            string `json:"llmModel"`
	EmbedModel          string `json:"embedModel"`
	ChunkMaxLength      string `json:"chunkMaxLength"`
	StartChunkThreshold string `json:"startChunkThreshold"`
	SimilarityThreshold string `json:"similarityThreshold"`
	EntriesPerFile      string `json:"entriesPerFile"`
}

type JobOutput interface {
	Output2()
}

var JobOutputTypeUnion = types.NewTypeUnion[JobOutput](
	(*AIJobOutput)(nil),
	(*BindingJobOutput)(nil),
)

var _ = serder.UseTypeUnionInternallyTagged(&JobOutputTypeUnion, "type")

type JobOutputBase struct{}

func (d *JobOutputBase) Output2() {}

type PublicOutput struct {
	serder.Metadata `union:"object"`
	JobOutputBase
	Type string `json:"type"`
}

type HPCOutput struct {
	serder.Metadata `union:"HPCSlurm"`
	JobOutputBase
	Type   string `json:"type"`
	Output string `json:"output"`
}

type AIJobOutput struct {
	serder.Metadata `union:"AI"`
	JobOutputBase
	Type   string `json:"type"`
	Output string `json:"output"`
}

type BindingJobOutput struct {
	serder.Metadata `union:"binding"`
	JobOutputBase
	Type      string `json:"type"`
	BindingID DataID `json:"bindingID"`
}

type DataReturnJobOutput struct {
	serder.Metadata `union:"DataReturn"`
	JobOutputBase
	Type          string               `json:"type"`
	ReportMessage TrainJobStatusReport `json:"report"`
	PackageID     cdssdk.PackageID     `json:"packageID"`
}

type JobStatusReport interface {
	Report()
}

var JobStatusReportTypeUnion = types.NewTypeUnion[JobStatusReport](
	(*TrainJobStatusReport)(nil),
	(*InferenceJobStatusReport)(nil),
)

var _ = serder.UseTypeUnionInternallyTagged(&JobStatusReportTypeUnion, "type")

type JobStatusReportBase struct{}

func (d *JobStatusReportBase) Report() {}

type TrainJobStatusReport struct {
	serder.Metadata `union:"Train"`
	JobStatusReportBase
	Type      string    `json:"type"`
	TaskName  string    `json:"taskName"`
	TaskID    string    `json:"taskID"`
	Status    bool      `json:"status"`
	Message   string    `json:"message"`
	ClusterID ClusterID `json:"clusterID"`
	Output    string    `json:"output"`
}

type InferenceJobStatusReport struct {
	serder.Metadata `union:"Inference"`
	JobStatusReportBase
	Type     string `json:"type"`
	TaskName string `json:"taskName"`
	TaskID   string `json:"taskID"`
	Status   bool   `json:"status"`
	Message  string `json:"message"`
	URL      string `json:"url"`
}