// Copyright 2025 The A2A Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package a2a

import (
	"encoding/gob"
	"encoding/json"
	"fmt"
	"time"

	"github.com/google/uuid"
)

// SendMessageResult represents a response for non-streaming message send.
type SendMessageResult interface {
	isSendMessageResult()
}

func (*Task) isSendMessageResult()    {}
func (*Message) isSendMessageResult() {}

// Event interface is used to represent types that can be sent over a streaming connection.
type Event interface {
	isEvent()
}

func (*Message) isEvent()                 {}
func (*Task) isEvent()                    {}
func (*TaskStatusUpdateEvent) isEvent()   {}
func (*TaskArtifactUpdateEvent) isEvent() {}

// MessageRole represents a set of possible values that identify the message sender.
type MessageRole string

const (
	MessageRoleAgent MessageRole = "agent"
	MessageRoleUser  MessageRole = "user"
)

// NewMessageID generates a new random message identifier.
func NewMessageID() string {
	return uuid.NewString()
}

// Message represents a single message in the conversation between a user and an agent.
type Message struct {
	// ID is a unique identifier for the message, typically a UUID, generated by the sender.
	ID string `json:"messageId" yaml:"messageId" mapstructure:"messageId"`

	// ContextID is the context identifier for this message, used to group related interactions.
	// An empty string means the message doesn't reference any context.
	ContextID string `json:"contextId,omitempty" yaml:"contextId,omitempty" mapstructure:"contextId,omitempty"`

	// Extensions are the URIs of extensions that are relevant to this message.
	Extensions []string `json:"extensions,omitempty" yaml:"extensions,omitempty" mapstructure:"extensions,omitempty"`

	// Metadata is an optional metadata for extensions. The key is an extension-specific identifier.
	Metadata map[string]any `json:"metadata,omitempty" yaml:"metadata,omitempty" mapstructure:"metadata,omitempty"`

	// Parts is an array of content parts that form the message body. A message can be
	// composed of multiple parts of different types (e.g., text and files).
	Parts ContentParts `json:"parts" yaml:"parts" mapstructure:"parts"`

	// ReferenceTasks is a list of other task IDs that this message references for additional context.
	ReferenceTasks []TaskID `json:"referenceTaskIds,omitempty" yaml:"referenceTaskIds,omitempty" mapstructure:"referenceTaskIds,omitempty"`

	// Role identifies the sender of the message.
	Role MessageRole `json:"role" yaml:"role" mapstructure:"role"`

	// TaskID is the identifier of the task this message is part of. Can be omitted for the
	// first message of a new task.
	// An empty string means the message doesn't reference any Task.
	TaskID TaskID `json:"taskId,omitempty" yaml:"taskId,omitempty" mapstructure:"taskId,omitempty"`
}

// NewMessage creates a new message with a random identifier.
func NewMessage(role MessageRole, parts ...Part) *Message {
	return &Message{
		ID:    NewMessageID(),
		Role:  role,
		Parts: parts,
	}
}

// NewMessageForTask creates a new message with a random identifier that references the provided Task.
func NewMessageForTask(role MessageRole, task *Task, parts ...Part) *Message {
	return &Message{
		ID:        NewMessageID(),
		Role:      role,
		TaskID:    task.ID,
		ContextID: task.ContextID,
		Parts:     parts,
	}
}

// TaskID is a unique identifier for the task, generated by the server for a new task.
type TaskID string

// NewTaskID creates a new random task identifier.
func NewTaskID() TaskID {
	return TaskID(uuid.NewString())
}

// NewContextID creates a new random context identifier.
func NewContextID() string {
	return uuid.NewString()
}

// TastState defines a set of possible task states.
type TaskState string

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

// Terminal returns true for states in which a Task becomes immutable, i.e. no further
// changes to the Task are permitted.
func (ts TaskState) Terminal() bool {
	return ts == TaskStateCompleted ||
		ts == TaskStateCanceled ||
		ts == TaskStateFailed ||
		ts == TaskStateRejected
}

// Task represents a single, stateful operation or conversation between a client and an agent.
type Task struct {
	// ID is a unique identifier for the task, generated by the server for a new task.
	ID TaskID `json:"id" yaml:"id" mapstructure:"id"`

	// Artifacts is a collection of artifacts generated by the agent during the execution of the task.
	Artifacts []*Artifact `json:"artifacts,omitempty" yaml:"artifacts,omitempty" mapstructure:"artifacts,omitempty"`

	// ContextID is a server-generated identifier for maintaining context across multiple related
	// tasks or interactions. Required to be non empty.
	ContextID string `json:"contextId" yaml:"contextId" mapstructure:"contextId"`

	// History is an array of messages exchanged during the task, representing the conversation history.
	History []*Message `json:"history,omitempty" yaml:"history,omitempty" mapstructure:"history,omitempty"`

	// Metadata is an optional metadata for extensions. The key is an extension-specific identifier.
	Metadata map[string]any `json:"metadata,omitempty" yaml:"metadata,omitempty" mapstructure:"metadata,omitempty"`

	// Status is the current status of the task, including its state and a descriptive message.
	Status TaskStatus `json:"status" yaml:"status" mapstructure:"status"`
}

// TaskStatus represents the status of a task at a specific point in time.
type TaskStatus struct {
	// Message is an optional, human-readable message providing more details about the current status.
	Message *Message `json:"message,omitempty" yaml:"message,omitempty" mapstructure:"message,omitempty"`

	// State is the current state of the task's lifecycle.
	State TaskState `json:"state" yaml:"state" mapstructure:"state"`

	// Timestamp is a datetime indicating when this status was recorded.
	Timestamp *time.Time `json:"timestamp,omitempty" yaml:"timestamp,omitempty" mapstructure:"timestamp,omitempty"`
}

// ArtifactID is a unique identifier for the artifact within the scope of the task.
type ArtifactID string

// NewArtifactID creates a new random artifact identifier.
func NewArtifactID() ArtifactID {
	return ArtifactID(uuid.NewString())
}

// Artifact represents a file, data structure, or other resource generated by an agent during a task.
type Artifact struct {
	// ID is a unique identifier for the artifact within the scope of the task.
	ID ArtifactID `json:"artifactId" yaml:"artifactId" mapstructure:"artifactId"`

	// Description is an optional, human-readable description of the artifact.
	Description string `json:"description,omitempty" yaml:"description,omitempty" mapstructure:"description,omitempty"`

	// Extensions are the URIs of extensions that are relevant to this artifact.
	Extensions []string `json:"extensions,omitempty" yaml:"extensions,omitempty" mapstructure:"extensions,omitempty"`

	// Metadata is an optional metadata for extensions. The key is an extension-specific identifier.
	Metadata map[string]any `json:"metadata,omitempty" yaml:"metadata,omitempty" mapstructure:"metadata,omitempty"`

	// Name is an optional, human-readable name for the artifact.
	Name string `json:"name,omitempty" yaml:"name,omitempty" mapstructure:"name,omitempty"`

	// Parts is an array of content parts that make up the artifact.
	Parts ContentParts `json:"parts" yaml:"parts" mapstructure:"parts"`
}

// TaskArtifactUpdateEvent is an event sent by the agent to notify the client that an artifact has been
// generated or updated. This is typically used in streaming models.
type TaskArtifactUpdateEvent struct {
	// Append indicates if the content of this artifact should be appended to a previously sent
	// artifact with the same ID.
	Append bool `json:"append,omitempty" yaml:"append,omitempty" mapstructure:"append,omitempty"`

	// Artifact is the artifact that was generated or updated.
	Artifact *Artifact `json:"artifact" yaml:"artifact" mapstructure:"artifact"`

	// ContextID is the context ID associated with the task. Required to be non-empty.
	ContextID string `json:"contextId" yaml:"contextId" mapstructure:"contextId"`

	// LastChunk indicates if this is the final chunk of the artifact.
	LastChunk bool `json:"lastChunk,omitempty" yaml:"lastChunk,omitempty" mapstructure:"lastChunk,omitempty"`

	// TaskID is the ID of the task this artifact belongs to.
	TaskID TaskID `json:"taskId" yaml:"taskId" mapstructure:"taskId"`

	// Metadata is an optional metadata for extensions.
	Metadata map[string]any `json:"metadata,omitempty" yaml:"metadata,omitempty" mapstructure:"metadata,omitempty"`
}

// NewArtifactEvent create a TaskArtifactUpdateEvent for an Artifact with a random ID.
func NewArtifactEvent(task *Task, parts ...Part) *TaskArtifactUpdateEvent {
	return &TaskArtifactUpdateEvent{
		ContextID: task.ContextID,
		TaskID:    task.ID,
		Artifact: &Artifact{
			ID:    NewArtifactID(),
			Parts: parts,
		},
	}
}

// NewArtifactUpdateEvent creates a TaskArtifactUpdateEvent that represents an update of the artifact with the provided ID.
func NewArtifactUpdateEvent(task *Task, id ArtifactID, parts ...Part) *TaskArtifactUpdateEvent {
	return &TaskArtifactUpdateEvent{
		ContextID: task.ContextID,
		TaskID:    task.ID,
		Append:    true,
		Artifact: &Artifact{
			ID:    id,
			Parts: parts,
		},
	}
}

// TaskStatusUpdateEvent is an event sent by the agent to notify the client of a change in a task's status.
// This is typically used in streaming or subscription models.
type TaskStatusUpdateEvent struct {
	// ContextID is the context ID associated with the task. Required to be non-empty.
	ContextID string `json:"contextId" yaml:"contextId" mapstructure:"contextId"`

	// Final indicates if this is the final event in the stream for this interaction.
	Final bool `json:"final" yaml:"final" mapstructure:"final"`

	// Status is the new status of the task.
	Status TaskStatus `json:"status" yaml:"status" mapstructure:"status"`

	// TaskID is the ID of the task that was updated.
	TaskID TaskID `json:"taskId" yaml:"taskId" mapstructure:"taskId"`

	// Metadata is an optional metadata for extensions.
	Metadata map[string]any `json:"metadata,omitempty" yaml:"metadata,omitempty" mapstructure:"metadata,omitempty"`
}

// NewStatusUpdateEvent creates a TaskStatusUpdateEvent that references the provided Task.
func NewStatusUpdateEvent(task *Task, state TaskState, msg *Message) *TaskStatusUpdateEvent {
	now := time.Now()
	return &TaskStatusUpdateEvent{
		ContextID: task.ContextID,
		TaskID:    task.ID,
		Status: TaskStatus{
			State:     state,
			Message:   msg,
			Timestamp: &now,
		},
	}
}

// ContentParts is an array of content parts that form the message body or an artifact.
type ContentParts []Part

func (j *ContentParts) UnmarshalJSON(b []byte) error {
	type typedPart struct {
		Kind string `json:"kind"`
	}

	var arr []json.RawMessage
	if err := json.Unmarshal(b, &arr); err != nil {
		return err
	}

	result := make([]Part, len(arr))
	for i, rawMsg := range arr {
		var tp typedPart
		if err := json.Unmarshal(rawMsg, &tp); err != nil {
			return err
		}
		switch tp.Kind {
		case "text":
			var part TextPart
			if err := json.Unmarshal(rawMsg, &part); err != nil {
				return err
			}
			result[i] = part
		case "data":
			var part DataPart
			if err := json.Unmarshal(rawMsg, &part); err != nil {
				return err
			}
			result[i] = part
		case "file":
			var part FilePart
			if err := json.Unmarshal(rawMsg, &part); err != nil {
				return err
			}
			result[i] = part
		default:
			return fmt.Errorf("unknown part kind %s", tp.Kind)
		}
	}

	*j = result
	return nil
}

// Part is a discriminated union representing a part of a message or artifact, which can
// be text, a file, or structured data.
type Part interface {
	isPart()

	Meta() map[string]any
}

func (TextPart) isPart() {}
func (FilePart) isPart() {}
func (DataPart) isPart() {}

func init() {
	gob.Register(TextPart{})
	gob.Register(FilePart{})
	gob.Register(DataPart{})
}

// TextPart represents a text segment within a message or artifact.
type TextPart struct {
	// Text is the string content of the text part.
	Text string `json:"text" yaml:"text" mapstructure:"text"`

	// Metadata is an optional metadata associated with the part.
	Metadata map[string]any `json:"metadata,omitempty" yaml:"metadata,omitempty" mapstructure:"metadata,omitempty"`
}

func (p TextPart) Meta() map[string]any {
	return p.Metadata
}

func (p TextPart) MarshalJSON() ([]byte, error) {
	type wrapped TextPart
	type withKind struct {
		Kind string `json:"kind"`
		wrapped
	}
	return json.Marshal(withKind{Kind: "text", wrapped: wrapped(p)})
}

// DataPart represents a structured data segment (e.g., JSON) within a message or artifact.
type DataPart struct {
	// Data is the structured data content.
	Data map[string]any `json:"data" yaml:"data" mapstructure:"data"`

	// Metadata is an optional metadata associated with the part.
	Metadata map[string]any `json:"metadata,omitempty" yaml:"metadata,omitempty" mapstructure:"metadata,omitempty"`
}

func (p DataPart) Meta() map[string]any {
	return p.Metadata
}

func (p DataPart) MarshalJSON() ([]byte, error) {
	type wrapped DataPart
	type withKind struct {
		Kind string `json:"kind"`
		wrapped
	}
	return json.Marshal(withKind{Kind: "data", wrapped: wrapped(p)})
}

// FilePart represents a file segment within a message or artifact. The file content can be
// provided either directly as bytes or as a URI.
type FilePart struct {
	// File the file content, represented as either a URI or as base64-encoded bytes.
	File FilePartContent `json:"file" yaml:"file" mapstructure:"file"`

	// Metadata is an optional metadata associated with the part.
	Metadata map[string]any `json:"metadata,omitempty" yaml:"metadata,omitempty" mapstructure:"metadata,omitempty"`
}

func (p FilePart) Meta() map[string]any {
	return p.Metadata
}

func (p FilePart) MarshalJSON() ([]byte, error) {
	type wrapped FilePart
	type withKind struct {
		Kind string `json:"kind"`
		wrapped
	}
	return json.Marshal(withKind{Kind: "file", wrapped: wrapped(p)})
}

func (p *FilePart) UnmarshalJSON(b []byte) error {
	type filePartContentUnion struct {
		FileMeta
		URI   string `json:"uri"`
		Bytes string `json:"bytes"`
	}
	type partJSON struct {
		File     filePartContentUnion `json:"file"`
		Metadata map[string]any       `json:"metadata"`
	}
	var decoded partJSON
	if err := json.Unmarshal(b, &decoded); err != nil {
		return err
	}

	if len(decoded.File.Bytes) == 0 && len(decoded.File.URI) == 0 {
		return fmt.Errorf("invalid file part: either Bytes or URI must be set")
	}
	if len(decoded.File.Bytes) > 0 && len(decoded.File.URI) > 0 {
		return fmt.Errorf("invalid file part: Bytes and URI cannot be set at the same time")
	}

	res := FilePart{Metadata: decoded.Metadata}
	if len(decoded.File.Bytes) > 0 {
		res.File = FileBytes{Bytes: decoded.File.Bytes, FileMeta: decoded.File.FileMeta}
	} else {
		res.File = FileURI{URI: decoded.File.URI, FileMeta: decoded.File.FileMeta}
	}

	*p = res
	return nil
}

// FilePartContent is a discriminated union of possible file part payloads.
type FilePartContent interface{ isFilePartContent() }

func (FileBytes) isFilePartContent() {}
func (FileURI) isFilePartContent()   {}

func init() {
	gob.Register(FileBytes{})
	gob.Register(FileURI{})
}

// FileMeta represents file metadata of a file part.
type FileMeta struct {
	// MimeType is an optinal MIME type of the file (e.g., "application/pdf").
	MimeType string `json:"mimeType,omitempty" yaml:"mimeType,omitempty" mapstructure:"mimeType,omitempty"`

	// Name is an optional name for the file (e.g., "document.pdf").
	Name string `json:"name,omitempty" yaml:"name,omitempty" mapstructure:"name,omitempty"`
}

// FileBytes represents a file with its content provided directly as a base64-encoded string.
type FileBytes struct {
	FileMeta
	// Bytes is the base64-encoded content of the file.
	Bytes string `json:"bytes" yaml:"bytes" mapstructure:"bytes"`
}

// FileURI represents a file with its content located at a specific URI.
type FileURI struct {
	FileMeta
	// URI is a URI pointing to the file's content.
	URI string `json:"uri" yaml:"uri" mapstructure:"uri"`
}

// Requests

// TaskIDParams defines parameters containing a task ID, used for simple task operations.
type TaskIDParams struct {
	// TaskID is the unique identifier of the task.
	ID TaskID `json:"id" yaml:"id" mapstructure:"id"`

	// Metadata is an optional metadata associated with the request.
	Metadata map[string]any `json:"metadata,omitempty" yaml:"metadata,omitempty" mapstructure:"metadata,omitempty"`
}

// TaskQueryParams defines parameters for querying a task, with an option to limit history length.
type TaskQueryParams struct {
	// HistoryLength is the number of most recent messages from the task's history to retrieve.
	HistoryLength *int `json:"historyLength,omitempty" yaml:"historyLength,omitempty" mapstructure:"historyLength,omitempty"`

	// ID is the unique identifier of the task.
	ID TaskID `json:"id" yaml:"id" mapstructure:"id"`

	// Metadata is an optional metadata associated with the request.
	Metadata map[string]any `json:"metadata,omitempty" yaml:"metadata,omitempty" mapstructure:"metadata,omitempty"`
}

// MessageSendConfig defines configuration options for a `message/send` or `message/stream` request.
type MessageSendConfig struct {
	// AcceptedOutputModes is a list of output MIME types the client is prepared to accept in the response.
	AcceptedOutputModes []string `json:"acceptedOutputModes,omitempty" yaml:"acceptedOutputModes,omitempty" mapstructure:"acceptedOutputModes,omitempty"`

	// Blocking indicates if the client will wait for the task to complete. The server may reject
	// this if the task is long-running.
	Blocking bool `json:"blocking,omitempty" yaml:"blocking,omitempty" mapstructure:"blocking,omitempty"`

	// HistoryLength is the number of most recent messages from the task's history to retrieve in the response.
	HistoryLength *int `json:"historyLength,omitempty" yaml:"historyLength,omitempty" mapstructure:"historyLength,omitempty"`

	// PushConfig is aonfiguration for the agent to send push notifications for updates after the initial response.
	PushConfig *PushConfig `json:"pushNotificationConfig,omitempty" yaml:"pushNotificationConfig,omitempty" mapstructure:"pushNotificationConfig,omitempty"`
}

// MessageSendParams defines the parameters for a request to send a message to an agent. This can be used
// to create a new task, continue an existing one, or restart a task.
type MessageSendParams struct {
	// Config is an optional configuration for the send request.
	Config *MessageSendConfig `json:"configuration,omitempty" yaml:"configuration,omitempty" mapstructure:"configuration,omitempty"`

	// Message is the message object being sent to the agent.
	Message *Message `json:"message" yaml:"message" mapstructure:"message"`

	// Metadata is an optional metadata for extensions.
	Metadata map[string]any `json:"metadata,omitempty" yaml:"metadata,omitempty" mapstructure:"metadata,omitempty"`
}
