package mlflow

import "encoding/json"

type ModelVersionStatus string

const ModelVersionStatusPendingRegistration = "PENDING_REGISTRATION"
const ModelVersionStatusFailedRegistration = "FAILED_REGISTRATION"
const ModelVersionStatusReady = "READY"

const ModelVersionStageNone = "None"
const ModelVersionStageStaging = "Staging"
const ModelVersionStageProduction = "Production"
const ModelVersionStageArchived = "Archived"

type RegisteredModelTag struct {
	Key   string `json:"key,omitempty"`
	Value string `json:"value,omitempty"`
}

type ModelVersionTag struct {
	Key   string `json:"key,omitempty"`
	Value string `json:"value,omitempty"`
}

type RegisteredModelAlias struct {
	Alias   string `json:"alias,omitempty"`
	Version string `json:"version,omitempty"`
}

type CreateRegisteredModelRequest struct {
	Name        string               `json:"name,omitempty"`
	Tags        []RegisteredModelTag `json:"tags,omitempty"`
	Description string               `json:"description,omitempty"`
}

func (r *CreateRegisteredModelRequest) serialize() []byte {
	rv, _ := json.Marshal(r)
	return rv
}

type CreateRegisteredModelResponse struct {
	RegisteredModel RegisteredModel `json:"registered_model,omitempty"`
}

type ModelVersion struct {
	Name                string             `json:"name,omitempty"`
	Version             string             `json:"version,omitempty"`
	CreationTimestamp   Timestamp          `json:"creation_timestamp,omitempty"`
	LastUpdateTimestamp Timestamp          `json:"last_update_timestamp,omitempty"`
	UserId              string             `json:"user_id,omitempty"`
	CurrentStage        string             `json:"current_stage,omitempty"`
	Description         string             `json:"description,omitempty"`
	Source              string             `json:"source,omitempty"`
	RunId               string             `json:"run_id,omitempty"`
	Status              ModelVersionStatus `json:"status,omitempty"`
	StatusMessage       string             `json:"status_message,omitempty"`
	Tags                []ModelVersionTag  `json:"tags,omitempty"`
	RunLink             string             `json:"run_link,omitempty"`
	Aliases             []string           `json:"aliases,omitempty"`
}

type RegisteredModel struct {
	Name                string                 `json:"name,omitempty"`
	CreationTimestamp   Timestamp              `json:"creation_timestamp,omitempty"`
	LastUpdateTimestamp Timestamp              `json:"last_update_timestamp,omitempty"`
	UserId              string                 `json:"user_id,omitempty"`
	Description         string                 `json:"description,omitempty"`
	LatestVersions      []ModelVersion         `json:"latest_versions,omitempty"`
	Tags                []RegisteredModelTag   `json:"tags,omitempty"`
	Aliases             []RegisteredModelAlias `json:"aliases,omitempty"`
}

type GetRegisteredModelResponse CreateRegisteredModelResponse

type RenameRegisteredModelRequrest struct {
	Name    string `json:"name,omitempty"`
	NewName string `json:"new_name,omitempty"`
}

func (r *RenameRegisteredModelRequrest) serialize() []byte {
	rv, _ := json.Marshal(r)
	return rv
}

type RenameRegisteredModelResponse CreateRegisteredModelResponse

type GetLatestModelVersionsRequest struct {
	Name   string   `json:"name,omitempty"`
	Stages []string `json:"stages,omitempty"`
}

func (r *GetLatestModelVersionsRequest) serialize() []byte {
	rv, _ := json.Marshal(r)
	return rv
}

type GetLatestModelVersionsReponse struct {
	ModelVersions []ModelVersion `json:"model_versions,omitempty"`
}

type GetModelVersionRequest struct {
	Name    string `json:"name,omitempty"`
	Version string `json:"version,omitempty"`
}

type GetModelVersionResponse struct {
	ModelVersion ModelVersion `json:"model_version,omitempty"`
}

type DeleteModelVersionRequest GetModelVersionRequest

func (r *DeleteModelVersionRequest) serialize() []byte {
	rv, _ := json.Marshal(r)
	return rv
}

type CreateModelVersionRequest struct {
	Name        string            `json:"name,omitempty"`
	Source      string            `json:"source,omitempty"`
	RunId       string            `json:"run_id,omitempty"`
	Tags        []ModelVersionTag `json:"tags,omitempty"`
	RunLink     string            `json:"run_link,omitempty"`
	Description string            `json:"description,omitempty"`
}

type CreateModelVersionResponse GetModelVersionResponse

func (r *CreateModelVersionRequest) serialize() []byte {
	rv, _ := json.Marshal(r)
	return rv
}

type GetDownloadUriForModelVersionArtifactsRequest GetModelVersionRequest

type GetDownloadUriForModelVersionArtifactsResponse struct {
	ArtifactUri string `json:"artifact_uri,omitempty"`
}

type SearchModelVersionsRequest struct {
	// Valid keys are '{'version_number', 'source_path', 'run_id', 'name'}'
	Filter    string `json:"filter,omitempty"`
	PageSize  int64  `json:"max_results,omitempty"`
	OrderBy   string `json:"order_by,omitempty"`
	PageToken string `json:"page_token,omitempty"`
}

func (e *SearchModelVersionsRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

func (e *SearchModelVersionsRequest) __init() *SearchModelVersionsRequest {
	if e.PageSize == 0 {
		e.PageSize = 1000
	}

	return e
}

type SearchModelVersionsResponse struct {
	ModelVersions []ModelVersion `json:"model_versions,omitempty"`
	NextPageToken string         `json:"next_page_token,omitempty"`
}

type UpdateModelVersionRequest struct {
	Name        string `json:"name,omitempty"`
	Version     string `json:"version,omitempty"`
	Description string `json:"description,omitempty"`
}

type UpdateModelVersionResponse GetModelVersionResponse

func (e *UpdateModelVersionRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

type UpdateRegisteredModelRequest struct {
	Name        string `json:"name,omitempty"`
	Description string `json:"description,omitempty"`
}

type UpdateRegisteredModelResponse GetRegisteredModelResponse

func (e *UpdateRegisteredModelRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

type TransitionModelVersionStageRequest struct {
	Name                    string `json:"name,omitempty"`
	Version                 string `json:"version,omitempty"`
	Stage                   string `json:"stage,omitempty"`
	ArchiveExistingVersions bool   `json:"archive_existing_versions,omitempty"`
}

func (e *TransitionModelVersionStageRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

type TransitionModelVersionStageResponse GetModelVersionResponse

type SearchRegisteredModelsRequest struct {
	// Valid keys are '{'name'}'
	Filter     string `json:"filter,omitempty"`
	MaxResults int64  `json:"max_results,omitempty"`
	OrderBy    string `json:"order_by,omitempty"`
	PageToken  string `json:"page_token,omitempty"`
}

func (e *SearchRegisteredModelsRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

func (e *SearchRegisteredModelsRequest) __init() *SearchRegisteredModelsRequest {
	if e.MaxResults == 0 {
		e.MaxResults = 1000
	}

	return e
}

type SearchRegisteredModelsResponse struct {
	RegisteredModels []RegisteredModel `json:"registered_models,omitempty"`
	NextPageToken    string            `json:"next_page_token,omitempty"`
}

type SetRegisteredModelTagRequest struct {
	RegisteredModelTag
	Name string `json:"name,omitempty"`
}

func (e *SetRegisteredModelTagRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

type SetModelVersionTagRequest struct {
	ModelVersionTag
	Name    string `json:"name,omitempty"`
	Version string `json:"version,omitempty"`
}

func (e *SetModelVersionTagRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

type DeleteRegisteredModelTagRequest struct {
	Key  string `json:"key,omitempty"`
	Name string `json:"name,omitempty"`
}

func (e *DeleteRegisteredModelTagRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

type DeleteModelVersionTagRequest struct {
	Key     string `json:"key,omitempty"`
	Name    string `json:"name,omitempty"`
	Version string `json:"version,omitempty"`
}

func (e *DeleteModelVersionTagRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

type SetRegisteredModelAliasRequest struct {
	Name    string `json:"name,omitempty"`
	Alias   string `json:"alias,omitempty"`
	Version string `json:"version,omitempty"`
}

func (e *SetRegisteredModelAliasRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

type GetModelVersionByAliasRequest struct {
	Name  string `json:"name,omitempty"`
	Alias string `json:"alias,omitempty"`
}

func (e *GetModelVersionByAliasRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

type DeleteRegisteredModelAliasRequest struct {
	Name  string `json:"name,omitempty"`
	Alias string `json:"alias,omitempty"`
}

func (e *DeleteRegisteredModelAliasRequest) serialize() []byte {
	j, _ := json.Marshal(e)
	return j
}

type GetModelVersionByAliasResponse GetModelVersionResponse

type SearchModelVersionsRequestAdvanced struct {
	SearchModelVersionsRequest
	Stage string
	Limit int64
}