package agent

import (
	"context"
	"encoding/json"
)

type Client interface {
	SendMessage(ctx context.Context, input SendMessageInput) (SendMessageResult, error)
	GetTask(ctx context.Context, input GetTaskInput) (Task, error)
	CancelTask(ctx context.Context, input CancelTaskInput) (Task, error)
}

type SendMessageInput struct {
	Text                string
	ContextID           string
	TaskID              string
	Metadata            map[string]any
	AcceptedOutputModes []string
	HistoryLength       int
	Blocking            bool
}

type SendMessageResult struct {
	Message *Message
	Task    *Task
	Raw     json.RawMessage
}

type GetTaskInput struct {
	TaskID        string
	HistoryLength int
	Metadata      map[string]any
}

type CancelTaskInput struct {
	TaskID   string
	Metadata map[string]any
}

type TaskState string

const (
	TaskStateSubmitted    TaskState = "submitted"
	TaskStateWorking      TaskState = "working"
	TaskStateInputNeeded  TaskState = "input-required"
	TaskStateCompleted    TaskState = "completed"
	TaskStateCanceled     TaskState = "canceled"
	TaskStateFailed       TaskState = "failed"
	TaskStateRejected     TaskState = "rejected"
	TaskStateAuthRequired TaskState = "auth-required"
	TaskStateUnknown      TaskState = "unknown"
)

type Task struct {
	Kind      string         `json:"kind,omitempty"`
	ID        string         `json:"id"`
	ContextID string         `json:"contextId,omitempty"`
	Status    TaskStatus     `json:"status"`
	Artifacts []Artifact     `json:"artifacts,omitempty"`
	History   []Message      `json:"history,omitempty"`
	Metadata  map[string]any `json:"metadata,omitempty"`
}

type TaskStatus struct {
	State     TaskState `json:"state"`
	Message   *Message  `json:"message,omitempty"`
	Timestamp string    `json:"timestamp,omitempty"`
}

type Message struct {
	Kind             string         `json:"kind,omitempty"`
	Role             string         `json:"role"`
	Parts            []Part         `json:"parts"`
	MessageID        string         `json:"messageId"`
	TaskID           string         `json:"taskId,omitempty"`
	ContextID        string         `json:"contextId,omitempty"`
	ReferenceTaskIDs []string       `json:"referenceTaskIds,omitempty"`
	Metadata         map[string]any `json:"metadata,omitempty"`
	Extensions       []string       `json:"extensions,omitempty"`
}

type Artifact struct {
	ArtifactID  string         `json:"artifactId,omitempty"`
	Name        string         `json:"name,omitempty"`
	Description string         `json:"description,omitempty"`
	Parts       []Part         `json:"parts,omitempty"`
	Metadata    map[string]any `json:"metadata,omitempty"`
}

type Part struct {
	Kind     string         `json:"kind"`
	Text     string         `json:"text,omitempty"`
	Data     map[string]any `json:"data,omitempty"`
	File     *FilePart      `json:"file,omitempty"`
	Metadata map[string]any `json:"metadata,omitempty"`
}

type FilePart struct {
	Name     string `json:"name,omitempty"`
	MimeType string `json:"mimeType,omitempty"`
	Bytes    string `json:"bytes,omitempty"`
	URI      string `json:"uri,omitempty"`
}