package mlflow

import (
	"bytes"
	"encoding/json"
	"time"
)

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

type CreateRunRequest struct {
	ExperimentId string      `json:"experiment_id,omitempty"`
	UserId       string      `json:"user_id,omitempty"`
	RunName      string      `json:"run_name,omitempty"`
	StartTime    TimestampMs `json:"start_time,omitempty"`
	Tags         []RunTag    `json:"tags,omitempty"`
}

type GetRunReply struct {
	Run Run `json:"run,omitempty"`
}

type CreateRunReply struct {
	GetRunReply
}

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

type RunStatus string

const RunStatusRunning = "RUNNING"
const RunStatusScheduled = "SCHEDULED"
const RunStatusFinished = "FINISHED"
const RunStatusFailed = "FAILED"
const RunStatusKilled = "KILLED"

type RunInfo struct {
	RunId          string      `json:"run_id,omitempty"`
	RunName        string      `json:"run_name,omitempty"`
	ExperimentId   string      `json:"experiment_id,omitempty"`
	UserId         string      `json:"user_id,omitempty"`
	Status         RunStatus   `json:"status,omitempty"`
	StartTime      TimestampMs `json:"start_time,omitempty"`
	EndTime        TimestampMs `json:"end_time,omitempty"`
	ArtifactUri    string      `json:"artifact_uri,omitempty"`
	LifecycleStage string      `json:"lifecycle_stage,omitempty"`
}

type Metric struct {
	Key       string      `json:"key,omitempty"`
	Value     float64     `json:"value,omitempty"`
	Timestamp TimestampMs `json:"timestamp,omitempty"`
	Step      int64       `json:"step,omitempty"`
}

func (r *Metric) __init() *Metric {
	if r.Timestamp == 0 {
		r.Timestamp = TimestampMs(time.Now().UnixMicro())
	}
	return r
}

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

type RunData struct {
	Metrics []Metric `json:"metics,omitempty"`
	Params  []Param  `json:"params,omitempty"`
	Tags    []RunTag `json:"tags,omitempty"`
}

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

type Dataset struct {
	Name       string `json:"name,omitempty"`
	Digest     string `json:"digest,omitempty"`
	SourceType string `json:"source_type,omitempty"`
	Source     string `json:"source,omitempty"`
	Schema     string `json:"schema,omitempty"`
	Profile    string `json:"profile,omitempty"`
}

type DatasetInput struct {
	Tags    []InputTag `json:"tags,omitempty"`
	Dataset Dataset    `json:"dataset,omitempty"`
}

type RunInputs struct {
	DatasetInputs []DatasetInput `json:"dataset_inputs,omitempty"`
}

type Run struct {
	Info   RunInfo   `json:"info,omitempty"`
	Data   RunData   `json:"data,omitempty"`
	Inputs RunInputs `json:"inputs,omitempty"`
}

type SearchRunRequest struct {
	ExperimentIds []string `json:"experiment_ids,omitempty"`
	Filter        string   `json:"filter,omitempty"`
	RunViewType   ViewType `json:"run_view_type,omitempty"`
	MaxResults    int32    `json:"max_results,omitempty"`
	OrgerBy       []string `json:"order_by,omitempty"`
	PageToken     string   `json:"page_token,omitempty"`
}

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

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

	return e
}

func (e *SearchRunRequest) getReader() *bytes.Reader {
	return bytes.NewReader(e.__init().serialize())
}

type LogMetricRequest struct {
	Metric
	RunId string `json:"run_id,omitempty"`
}

func (fe *LogMetricRequest) __init() *LogMetricRequest { fe.Metric.__init(); return fe }

type LogBatchRequest struct {
	RunId   string   `json:"run_id,omitempty"`
	Metrics []Metric `json:"metrics,omitempty"`
	Params  []Param  `json:"params,omitempty"`
	Tags    []RunTag `json:"tags,omitempty"`
}

type LogInputsRequest struct {
	RunId    string         `json:"run_id,omitempty"`
	Datasets []DatasetInput `json:"datasets,omitempty"`
}

type SetTagRequest struct {
	RunTag
	RunId string `json:"run_id,omitempty"`
}

type DeleteTagRequest struct {
	RunTag
	RunId string `json:"run_id,omitempty"`
}

type LogParamRequest struct {
	Param
	RunId string `json:"run_id,omitempty"`
}

type GetMetricHistoryResponse struct {
	Metrics       []Metric `json:"metrics,omitempty"`
	NextPageToken string   `json:"nex_page_token,omitempty"`
}

type GetMetricHistoryRequest struct {
	RunId      string `json:"run_id,omitempty"`
	MetricKey  string `json:"metric_key,omitempty"`
	PageToken  string `json:"page_token,omitempty"`
	MaxResults int32  `json:"max_results,omitempty"`
}

func (r *GetMetricHistoryRequest) __init() *GetMetricHistoryRequest {
	if r.MaxResults == 0 {
		r.MaxResults = 1000
	}
	return r
}