sdk

type Config struct {
	ManagerURL      string
	TLSVerification bool
}

type ExperimentConfig struct {
	ExperimentID  string         `json:"experiment_id"`
	RoundID       string         `json:"round_id"`
	ModelRef      string         `json:"model_ref"`
	Participants  []string       `json:"participants"`
	Hyperparams   map[string]any `json:"hyperparams"`
	KOfN          int            `json:"k_of_n"`
	TimeoutS      int            `json:"timeout_s"`
	TaskWasmImage string         `json:"task_wasm_image,omitempty"`
}

type ExperimentResult struct {
	ExperimentID string `json:"experiment_id"`
	RoundID      string `json:"round_id"`
	Status       string `json:"status"`
}

type FLTask struct {
	RoundID     string         `json:"round_id"`
	ModelRef    string         `json:"model_ref"`
	Config      map[string]any `json:"config"`
	Hyperparams map[string]any `json:"hyperparams,omitempty"`
}

type FLUpdate struct {
	RoundID      string         `json:"round_id"`
	PropletID    string         `json:"proplet_id"`
	BaseModelURI string         `json:"base_model_uri"`
	NumSamples   int            `json:"num_samples"`
	Metrics      map[string]any `json:"metrics"`
	Update       map[string]any `json:"update"`
	ReceivedAt   time.Time      `json:"received_at"`
}

type HealthInfo struct {
	Status      string `json:"status"`
	Version     string `json:"version"`
	Commit      string `json:"commit"`
	Description string `json:"description"`
	BuildTime   string `json:"build_time"`
	InstanceID  string `json:"instance_id"`
}

type JobPage struct {
	Offset uint64       `json:"offset"`
	Limit  uint64       `json:"limit"`
	Total  uint64       `json:"total"`
	Jobs   []JobSummary `json:"jobs"`
}

type JobRequest struct {
	Name          string `json:"name"`
	Tasks         []Task `json:"tasks"`
	ExecutionMode string `json:"execution_mode,omitempty"`
}

type JobResponse struct {
	JobID string `json:"job_id"`
	Tasks []Task `json:"tasks"`
}

type JobSummary struct {
	JobID      string    `json:"job_id"`
	Name       string    `json:"name,omitempty"`
	State      uint8     `json:"state"`
	Tasks      []Task    `json:"tasks"`
	StartTime  time.Time `json:"start_time"`
	FinishTime time.Time `json:"finish_time"`
	CreatedAt  time.Time `json:"created_at"`
}

type PageMetadata struct {
	Offset   uint64        `json:"offset"`
	Limit    uint64        `json:"limit"`
	Metadata task.Metadata `json:"metadata,omitempty"`
}

type Proplet struct {
	ID          string     `json:"id"`
	Name        string     `json:"name"`
	TaskCount   uint64     `json:"task_count"`
	Alive       bool       `json:"alive"`
	LastAliveAt *time.Time `json:"last_alive_at,omitempty"`
}

type PropletMetrics struct {
	PropletID string                `json:"proplet_id"`
	Namespace string                `json:"namespace"`
	Timestamp time.Time             `json:"timestamp"`
	CPU       proplet.CPUMetrics    `json:"cpu_metrics"`
	Memory    proplet.MemoryMetrics `json:"memory_metrics"`
}

type PropletMetricsPage struct {
	Offset  uint64           `json:"offset"`
	Limit   uint64           `json:"limit"`
	Total   uint64           `json:"total"`
	Metrics []PropletMetrics `json:"metrics"`
}

type PropletPage struct {
	Offset   uint64    `json:"offset"`
	Limit    uint64    `json:"limit"`
	Total    uint64    `json:"total"`
	Proplets []Proplet `json:"proplets"`
}

type RoundStatus struct {
	RoundID      string `json:"round_id"`
	Completed    bool   `json:"completed"`
	NumUpdates   int    `json:"num_updates"`
	KOfN         int    `json:"k_of_n"`
	ModelVersion int    `json:"model_version,omitempty"`
}

type SDK interface {
	// CreateTask creates a new task.
	//
	// example:
	//  task := sdk.Task{
	//    Name:	 "John Doe"
	//  }
	//  task, _ := sdk.CreateTask(task)
	//  fmt.Println(task)
	CreateTask(task Task) (Task, error)

	// GetTask gets a task by id.
	//
	// example:
	//  task, _ := sdk.GetTask("b1d10738-c5d7-4ff1-8f4d-b9328ce6f040")
	//  fmt.Println(task)
	GetTask(id string) (Task, error)

	// ListTasks lists tasks with optional metadata filtering.
	//
	// example:
	//  taskPage, _ := sdk.ListTasks(sdk.PageMetadata{Offset: 0, Limit: 10})
	//  fmt.Println(taskPage)
	ListTasks(pm PageMetadata) (TaskPage, error)

	// UpdateTask updates a task.
	//
	// example:
	//  task := sdk.Task{
	//    Name:	 "John Doe"
	//  }
	//  task, _ := sdk.UpdateTask(task)
	//  fmt.Println(task)
	UpdateTask(task Task) (Task, error)

	// UploadTaskFile uploads a Wasm binary for a task via multipart form.
	//
	// example:
	//  task, _ := sdk.UploadTaskFile("b1d10738-c5d7-4ff1-8f4d-b9328ce6f040", "/path/to/app.wasm")
	//  fmt.Println(task)
	UploadTaskFile(id string, filePath string) (Task, error)

	// GetTaskResults returns the stored execution results of a task.
	//
	// example:
	//  results, _ := sdk.GetTaskResults("b1d10738-c5d7-4ff1-8f4d-b9328ce6f040")
	//  fmt.Println(results)
	GetTaskResults(id string) (any, error)

	// GetTaskMetrics returns the paginated metrics for a task.
	//
	// example:
	//  page, _ := sdk.GetTaskMetrics("b1d10738-c5d7-4ff1-8f4d-b9328ce6f040", 0, 10)
	//  fmt.Println(page)
	GetTaskMetrics(id string, offset, limit uint64) (TaskMetricsPage, error)

	// DeleteTask deletes a task.
	//
	// example:
	//  task, _ := sdk.DeleteTask("b1d10738-c5d7-4ff1-8f4d-b9328ce6f040")
	//  fmt.Println(task)
	DeleteTask(id string) error

	// StartTask starts a task.
	//
	// example:
	//  task, _ := sdk.StartTask("b1d10738-c5d7-4ff1-8f4d-b9328ce6f040")
	//  fmt.Println(task)
	StartTask(id string) error

	// StopTask stops a task.
	//
	// example:
	//  task, _ := sdk.StopTask("b1d10738-c5d7-4ff1-8f4d-b9328ce6f040")
	//  fmt.Println(task)
	StopTask(id string) error

	InvokeTask(id string, inputs []string, env map[string]string) (string, error)

	// CreateJob creates a new job with multiple tasks.
	//
	// example:
	//  req := sdk.JobRequest{
	//    Name: "my-job",
	//    Tasks: []sdk.Task{...},
	//    ExecutionMode: "parallel",
	//  }
	//  job, _ := sdk.CreateJob(req)
	CreateJob(req JobRequest) (JobResponse, error)

	// GetJob gets a job by id.
	//
	// example:
	//  job, _ := sdk.GetJob("b1d10738-c5d7-4ff1-8f4d-b9328ce6f040")
	GetJob(jobID string) (JobResponse, error)

	// ListJobs lists jobs with optional status filter.
	// Status can be "pending", "running", "completed", "failed", or "" (all).
	//
	// example:
	//  jobPage, _ := sdk.ListJobs(0, 10, "")
	//  jobPage, _ := sdk.ListJobs(0, 10, "running")
	ListJobs(offset uint64, limit uint64, status string) (JobPage, error)

	// StartJob starts a job.
	//
	// example:
	//  _ := sdk.StartJob("b1d10738-c5d7-4ff1-8f4d-b9328ce6f040")
	StartJob(jobID string) error

	// StopJob stops a job.
	//
	// example:
	//  _ := sdk.StopJob("b1d10738-c5d7-4ff1-8f4d-b9328ce6f040")
	StopJob(jobID string) error

	// GetProplet returns a single proplet by id.
	//
	// example:
	//  p, _ := sdk.GetProplet("b1d10738-c5d7-4ff1-8f4d-b9328ce6f040")
	//  fmt.Println(p)
	GetProplet(id string) (Proplet, error)

	// GetPropletAliveHistory returns the paginated heartbeat history for a proplet.
	//
	// example:
	//  page, _ := sdk.GetPropletAliveHistory("b1d10738-c5d7-4ff1-8f4d-b9328ce6f040", 0, 10)
	//  fmt.Println(page)
	GetPropletAliveHistory(id string, offset, limit uint64) (proplet.PropletAliveHistoryPage, error)

	// GetPropletMetrics returns the paginated metrics for a proplet.
	//
	// example:
	//  page, _ := sdk.GetPropletMetrics("b1d10738-c5d7-4ff1-8f4d-b9328ce6f040", 0, 10)
	//  fmt.Println(page)
	GetPropletMetrics(id string, offset, limit uint64) (PropletMetricsPage, error)

	// ListProplets returns a paginated list of proplets, optionally filtered by status.
	//
	// example:
	//  page, _ := sdk.ListProplets(0, 10, "")
	//  fmt.Println(page)
	ListProplets(offset, limit uint64, status string) (PropletPage, error)

	// GetPropletSDF returns the SDF description of a proplet.
	//
	// example:
	//  doc, _ := sdk.GetPropletSDF("b1d10738-c5d7-4ff1-8f4d-b9328ce6f040")
	//  fmt.Println(doc)
	GetPropletSDF(id string) (sdf.Document, error)

	// DeleteProplet deletes a proplet by id.
	//
	// example:
	//  err := sdk.DeleteProplet("b1d10738-c5d7-4ff1-8f4d-b9328ce6f040")
	//  fmt.Println(err)
	DeleteProplet(id string) error

	// CreateWorkflow creates a multi-task workflow (DAG).
	//
	// example:
	//  tasks, _ := sdk.CreateWorkflow([]sdk.Task{
	//    {Name: "step-1"},
	//    {Name: "step-2", DependsOn: []string{"<step-1-id>"}},
	//  })
	//  fmt.Println(tasks)
	CreateWorkflow(tasks []Task) ([]Task, error)

	// ConfigureExperiment configures a federated learning experiment.
	//
	// example:
	//  result, _ := sdk.ConfigureExperiment(sdk.ExperimentConfig{
	//    ExperimentID: "exp-001",
	//    RoundID: "round-1",
	//  })
	//  fmt.Println(result)
	ConfigureExperiment(config ExperimentConfig) (ExperimentResult, error)

	// GetFLTask returns the federated learning task for the current round.
	//
	// example:
	//  task, _ := sdk.GetFLTask("round-1", "proplet-1")
	//  fmt.Println(task)
	GetFLTask(roundID, propletID string) (FLTask, error)

	// PostFLUpdate submits a model update in JSON format.
	//
	// example:
	//  err := sdk.PostFLUpdate(sdk.FLUpdate{RoundID: "round-1"})
	PostFLUpdate(update FLUpdate) error

	// PostFLUpdateCBOR submits a model update in CBOR format.
	//
	// example:
	//  err := sdk.PostFLUpdateCBOR(data)
	PostFLUpdateCBOR(data []byte) error

	// GetRoundStatus returns the completion status of a federated learning round.
	//
	// example:
	//  status, _ := sdk.GetRoundStatus("round-1")
	//  fmt.Println(status)
	GetRoundStatus(roundID string) (RoundStatus, error)

	// GetHealth returns the manager health status.
	//
	// example:
	//  info, _ := sdk.GetHealth()
	//  fmt.Println(info)
	GetHealth() (HealthInfo, error)
}

type Task struct {
	ID       string            `json:"id,omitempty"`
	Name     string            `json:"name"`
	Kind     string            `json:"kind,omitempty"`
	State    uint8             `json:"state,omitempty"`
	Mode     string            `json:"mode,omitempty"`
	ImageURL string            `json:"image_url,omitempty"`
	JobID    string            `json:"job_id,omitempty"`
	CLIArgs  []string          `json:"cli_args,omitempty"`
	Env      map[string]string `json:"env,omitempty"`
	// File is the Wasm binary. When sending, set it to the base64-encoded
	// binary (format: byte). When reading task responses the manager redacts
	// the payload (e.g. "AGFzbQEAAA<REDACTED>RpdmFsdWU="), so a plain string is
	// used instead of []byte to avoid a base64 decoding failure.
	File            string         `json:"file,omitempty"`
	Inputs          []string       `json:"inputs,omitempty"`
	Daemon          bool           `json:"daemon,omitempty"`
	Latent          bool           `json:"latent,omitempty"`
	Encrypted       bool           `json:"encrypted,omitempty"`
	KBSResourcePath string         `json:"kbs_resource_path,omitempty"`
	PropletID       string         `json:"proplet_id,omitempty"`
	DependsOn       []string       `json:"depends_on,omitempty"`
	RunIf           string         `json:"run_if,omitempty"`
	WorkflowID      string         `json:"workflow_id,omitempty"`
	Broadcast       bool           `json:"broadcast,omitempty"`
	Priority        int            `json:"priority,omitempty"`
	Schedule        string         `json:"schedule,omitempty"`
	Timezone        string         `json:"timezone,omitempty"`
	IsRecurring     bool           `json:"is_recurring,omitempty"`
	Metadata        map[string]any `json:"metadata,omitempty"`
	StartTime       time.Time      `json:"start_time"`
	FinishTime      time.Time      `json:"finish_time"`
	CreatedAt       time.Time      `json:"created_at"`
	UpdatedAt       time.Time      `json:"updated_at"`
	Results         any            `json:"results,omitempty"`
}

type TaskMetrics struct {
	TaskID     string                     `json:"task_id"`
	PropletID  string                     `json:"proplet_id"`
	Metrics    proplet.ProcessMetrics     `json:"metrics"`
	Aggregated *proplet.AggregatedMetrics `json:"aggregated,omitempty"`
	Timestamp  time.Time                  `json:"timestamp"`
}

type TaskMetricsPage struct {
	Offset  uint64        `json:"offset"`
	Limit   uint64        `json:"limit"`
	Total   uint64        `json:"total"`
	Metrics []TaskMetrics `json:"metrics"`
}

type TaskPage struct {
	Offset uint64 `json:"offset"`
	Limit  uint64 `json:"limit"`
	Total  uint64 `json:"total"`
	Tasks  []Task `json:"tasks"`
}