// 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 taskupdate

import (
	"context"
	"fmt"

	"github.com/a2aproject/a2a-go/a2a"
)

// Saver is used for saving the Task after updating its state.
type Saver interface {
	Save(ctx context.Context, task *a2a.Task) error
}

// Manager is used for processing a2a.Event related to a Task. It updates
// the Task accordingly and uses Saver to store the new state.
type Manager struct {
	task  *a2a.Task
	saver Saver
}

// NewManager creates an initialized update Manager for the provided task.
func NewManager(saver Saver, task *a2a.Task) *Manager {
	return &Manager{task: task, saver: saver}
}

// Process validates that the event is associated with the managed Task and updates the Task accordingly.
func (mgr *Manager) Process(ctx context.Context, event a2a.Event) (*a2a.Task, error) {
	if mgr.task == nil {
		return nil, fmt.Errorf("event processor Task not set")
	}

	switch v := event.(type) {
	case *a2a.Message:
		return mgr.task, nil

	case *a2a.Task:
		if err := mgr.validate(v.ID, v.ContextID); err != nil {
			return nil, err
		}
		if err := mgr.saver.Save(ctx, v); err != nil {
			return nil, err
		}
		mgr.task = v
		return v, nil

	case *a2a.TaskArtifactUpdateEvent:
		if err := mgr.validate(v.TaskID, v.ContextID); err != nil {
			return nil, err
		}
		return mgr.updateArtifact(ctx, v)

	case *a2a.TaskStatusUpdateEvent:
		if err := mgr.validate(v.TaskID, v.ContextID); err != nil {
			return nil, err
		}
		return mgr.updateStatus(ctx, v)

	default:
		return nil, fmt.Errorf("unexpected event type %T", v)
	}
}

func (mgr *Manager) updateArtifact(_ context.Context, _ *a2a.TaskArtifactUpdateEvent) (*a2a.Task, error) {
	return nil, fmt.Errorf("not implemented")
}

func (mgr *Manager) updateStatus(ctx context.Context, event *a2a.TaskStatusUpdateEvent) (*a2a.Task, error) {
	task := mgr.task

	if task.Status.Message != nil {
		task.History = append(task.History, task.Status.Message)
	}

	if event.Metadata != nil {
		if task.Metadata == nil {
			task.Metadata = make(map[string]any)
		}
		for k, v := range event.Metadata {
			task.Metadata[k] = v
		}
	}

	task.Status = event.Status

	if err := mgr.saver.Save(ctx, task); err != nil {
		return nil, err
	}

	return task, nil
}

func (mgr *Manager) validate(taskID a2a.TaskID, contextID string) error {
	if mgr.task.ID != taskID {
		return fmt.Errorf("task IDs don't match: %s != %s", mgr.task.ID, taskID)
	}

	if mgr.task.ContextID != contextID {
		return fmt.Errorf("context IDs don't match: %s != %s", mgr.task.ContextID, contextID)
	}

	return nil
}
