package p2p

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"sync"
	"time"

	"synapse/types"
)

// =============================================================================
// P2P Protocol Handler
// =============================================================================

// TaskHandler is the interface for handling task execution
type TaskHandler interface {
	// ExecuteTask executes a task and returns the result
	ExecuteTask(ctx context.Context, taskType string, payload map[string]interface{}) (any, error)

	// ValidateTask validates if the task can be executed
	ValidateTask(taskType string, payload map[string]interface{}) ([]types.MissingField, bool)

	// GetCapabilities returns the agent's capabilities
	GetCapabilities() []string
}

// MessageHandler is the interface for handling incoming messages
type MessageHandler interface {
	// HandleMessage handles an incoming P2P message
	HandleMessage(msg *types.UnifiedSynapseMessage) (*types.UnifiedSynapseMessage, error)
}

// CallbackSender is the interface for sending callbacks to remote agents
type CallbackSender interface {
	// SendCallback sends a message to a remote agent
	SendCallback(ctx context.Context, endpoint string, msg *types.UnifiedSynapseMessage) error
}

// =============================================================================
// P2PProtocol represents the P2P protocol handler
// =============================================================================

// P2PProtocol handles point-to-point Synapse communication
type P2PProtocol struct {
	mu sync.RWMutex

	// Agent identity
	agentEndpoint types.AgentEndpoint

	// Task state
	tasks map[string]*types.P2PTaskStatus

	// Handlers
	taskHandler    TaskHandler
	callbackSender CallbackSender

	// Configuration
	config *P2PConfig
}

// P2PConfig holds P2P protocol configuration
type P2PConfig struct {
	MaxRetries        int           // Max callback retries
	RetryBaseDelay    time.Duration // Base delay for retries
	TaskTimeout       time.Duration // Task execution timeout
	HeartbeatInterval time.Duration // Heartbeat interval
}

// DefaultP2PConfig returns default P2P configuration
func DefaultP2PConfig() *P2PConfig {
	return &P2PConfig{
		MaxRetries:        5,
		RetryBaseDelay:    2 * time.Second,
		TaskTimeout:       5 * time.Minute,
		HeartbeatInterval: 30 * time.Second,
	}
}

// =============================================================================
// P2PTask represents an active task
// =============================================================================

type P2PTask struct {
	ID            string
	Type          string
	Priority      types.Priority
	Status        types.P2PTaskStatus
	Payload       map[string]interface{}
	Sender        types.AgentEndpoint
	Receiver      types.AgentEndpoint
	CorrelationID string
	CreatedAt     time.Time
	UpdatedAt     time.Time
	CompletedAt   *time.Time
	Result        any
	Error         *types.SynapseError
}

// =============================================================================
// NewP2PProtocol creates a new P2P protocol handler
// =============================================================================

func NewP2PProtocol(agentEndpoint types.AgentEndpoint, config *P2PConfig) *P2PProtocol {
	if config == nil {
		config = DefaultP2PConfig()
	}
	return &P2PProtocol{
		agentEndpoint: agentEndpoint,
		tasks:         make(map[string]*types.P2PTaskStatus),
		config:        config,
	}
}

// SetTaskHandler sets the task handler
func (p *P2PProtocol) SetTaskHandler(h TaskHandler) {
	p.taskHandler = h
}

// SetCallbackSender sets the callback sender
func (p *P2PProtocol) SetCallbackSender(s CallbackSender) {
	p.callbackSender = s
}

// AgentEndpoint returns the agent's endpoint
func (p *P2PProtocol) AgentEndpoint() types.AgentEndpoint {
	return p.agentEndpoint
}

// =============================================================================
// HandleMessage processes an incoming P2P message
// =============================================================================

func (p *P2PProtocol) HandleMessage(ctx context.Context, msg *types.UnifiedSynapseMessage) (*types.UnifiedSynapseMessage, error) {
	p.mu.Lock()
	defer p.mu.Unlock()

	// Validate message
	if err := p.validateMessage(msg); err != nil {
		a2aErr, ok := err.(*types.SynapseError)
		if !ok {
			a2aErr = types.NewSynapseError(types.ErrMessageFormatInvalid, err.Error(), nil)
		}
		return p.newNACK(msg, a2aErr)
	}

	// Route based on message type
	switch types.P2PMessageType(msg.Type) {
	case types.P2PTaskRequest:
		return p.handleTaskRequest(ctx, msg)
	case types.P2PTaskConfirm:
		return p.handleTaskConfirm(ctx, msg)
	case types.P2PTaskReject:
		return p.handleTaskReject(ctx, msg)
	case types.P2PTaskNeedInfo:
		return p.handleTaskNeedInfo(ctx, msg)
	case types.P2PTaskInfoProvided:
		return p.handleTaskInfoProvided(ctx, msg)
	case types.P2PTaskStatusUpdate:
		return p.handleTaskStatusUpdate(ctx, msg)
	case types.P2PTaskComplete:
		return p.handleTaskComplete(ctx, msg)
	case types.P2PTaskFailed:
		return p.handleTaskFailed(ctx, msg)
	case types.P2PTaskCancel:
		return p.handleTaskCancel(ctx, msg)
	case types.P2PTaskAck:
		return p.handleTaskAck(ctx, msg)
	case types.P2PTaskNeedHuman:
		return p.handleTaskNeedHuman(ctx, msg)
	case types.P2PHumanDecision:
		return p.handleHumanDecision(ctx, msg)
	case types.P2PAck:
		return p.handleAck(ctx, msg)
	case types.P2PNack:
		return p.handleNack(ctx, msg)
	case types.P2PHeartbeat:
		return p.handleHeartbeat(ctx, msg)
	default:
		return p.newNACK(msg, types.NewSynapseError(
			types.ErrMessageFormatInvalid,
			fmt.Sprintf("Unknown message type: %s", msg.Type),
			nil,
		))
	}
}

// =============================================================================
// Message Handlers
// =============================================================================

// handleTaskRequest handles incoming task requests
func (p *P2PProtocol) handleTaskRequest(ctx context.Context, msg *types.UnifiedSynapseMessage) (*types.UnifiedSynapseMessage, error) {
	var payload types.TaskRequestPayload
	if err := json.Unmarshal(msg.Payload, &payload); err != nil {
		return p.newNACK(msg, types.NewSynapseError(
			types.ErrMessageFormatInvalid,
			"Invalid task request payload",
			err.Error(),
		))
	}

	taskID := payload.Task.ID
	log.Printf("[P2P] Task request received: %s (type: %s)", taskID, payload.Task.Type)

	// Check if we have a task handler
	if p.taskHandler == nil {
		return p.newReject(msg, taskID, "agent_not_ready", "Task handler not configured", nil)
	}

	// Validate the task
	missingFields, canHandle := p.taskHandler.ValidateTask(payload.Task.Type, payload.Task.Payload)
	if !canHandle {
		if len(missingFields) > 0 {
			// Need more info
			return p.newNeedInfo(msg, taskID, missingFields)
		}
		return p.newReject(msg, taskID, "capability_lacking", "Agent cannot handle this task type", nil)
	}

	// Create task context
	taskCtx, cancel := context.WithTimeout(ctx, p.config.TaskTimeout)
	defer cancel()

	// Execute task asynchronously
	go p.executeTask(taskCtx, msg, taskID, payload.Task.Type, payload.Task.Payload)

	// Return immediate confirmation
	return p.newAck(msg, taskID, "task_accepted")
}

// executeTask executes a task and handles callbacks
func (p *P2PProtocol) executeTask(ctx context.Context, origMsg *types.UnifiedSynapseMessage, taskID, taskType string, payload map[string]interface{}) {
	// Send initial status update
	p.sendStatusUpdate(ctx, taskID, types.P2PStatusExecuting, 0, "Task started")

	// Execute the task
	result, err := p.taskHandler.ExecuteTask(ctx, taskType, payload)

	if err != nil {
		// Task failed
		p.sendTaskFailed(ctx, taskID, "execution_failed", err.Error(), true, nil)
		return
	}

	// Task completed
	p.sendTaskComplete(ctx, taskID, result)
}

// sendStatusUpdate sends a status update to the sender
func (p *P2PProtocol) sendStatusUpdate(ctx context.Context, taskID string, status types.P2PTaskStatus, progress int, message string) {
	if p.callbackSender == nil {
		return
	}

	payload := types.TaskStatusUpdatePayload{
		TaskID:   taskID,
		Status:   string(status),
		Progress: progress,
		Message:  message,
	}

	payloadBytes, _ := json.Marshal(payload)

	response := &types.UnifiedSynapseMessage{
		ID:        types.GenerateMessageID(),
		Version:   "1.0.0",
		Type:      string(types.P2PTaskStatusUpdate),
		Sender:    p.agentEndpoint,
		Payload:   payloadBytes,
		Timestamp: types.NowTimestamp(),
	}

	if err := p.callbackSender.SendCallback(ctx, "", response); err != nil {
		log.Printf("[P2P] Failed to send status update: %v", err)
	}
}

// sendTaskComplete sends task completion notification
func (p *P2PProtocol) sendTaskComplete(ctx context.Context, taskID string, result any) {
	if p.callbackSender == nil {
		return
	}

	payload := types.TaskCompletePayload{}
	payload.TaskID = taskID
	payload.Result.OutputType = "task_result"
	payload.Result.Output = result

	payloadBytes, _ := json.Marshal(payload)

	response := &types.UnifiedSynapseMessage{
		ID:        types.GenerateMessageID(),
		Version:   "1.0.0",
		Type:      string(types.P2PTaskComplete),
		Sender:    p.agentEndpoint,
		Payload:   payloadBytes,
		Timestamp: types.NowTimestamp(),
	}

	if err := p.callbackSender.SendCallback(ctx, "", response); err != nil {
		log.Printf("[P2P] Failed to send task complete: %v", err)
	}
}

// sendTaskFailed sends task failure notification
func (p *P2PProtocol) sendTaskFailed(ctx context.Context, taskID, errorCode, errorMessage string, recoverable bool, partialResult any) {
	if p.callbackSender == nil {
		return
	}

	payload := types.TaskFailedPayload{}
	payload.TaskID = taskID
	payload.Error.Code = errorCode
	payload.Error.Message = errorMessage
	payload.Error.Recoverable = recoverable
	payload.PartialResult = partialResult

	payloadBytes, _ := json.Marshal(payload)

	response := &types.UnifiedSynapseMessage{
		ID:        types.GenerateMessageID(),
		Version:   "1.0.0",
		Type:      string(types.P2PTaskFailed),
		Sender:    p.agentEndpoint,
		Payload:   payloadBytes,
		Timestamp: types.NowTimestamp(),
	}

	if err := p.callbackSender.SendCallback(ctx, "", response); err != nil {
		log.Printf("[P2P] Failed to send task failed: %v", err)
	}
}

// handleTaskConfirm handles task confirmation from receiver
func (p *P2PProtocol) handleTaskConfirm(ctx context.Context, msg *types.UnifiedSynapseMessage) (*types.UnifiedSynapseMessage, error) {
	var payload types.TaskConfirmPayload
	if err := json.Unmarshal(msg.Payload, &payload); err != nil {
		return p.newNACK(msg, types.NewSynapseError(types.ErrMessageFormatInvalid, "Invalid confirm payload", err.Error()))
	}

	log.Printf("[P2P] Task confirmed: %s (estimated: %ds)", payload.TaskID, payload.EstimatedDuration)
	return p.newAck(msg, payload.TaskID, "confirmed")
}

// handleTaskReject handles task rejection from receiver
func (p *P2PProtocol) handleTaskReject(ctx context.Context, msg *types.UnifiedSynapseMessage) (*types.UnifiedSynapseMessage, error) {
	var payload types.TaskRejectPayload
	if err := json.Unmarshal(msg.Payload, &payload); err != nil {
		return p.newNACK(msg, types.NewSynapseError(types.ErrMessageFormatInvalid, "Invalid reject payload", err.Error()))
	}

	log.Printf("[P2P] Task rejected: %s (%s: %s)", payload.TaskID, payload.Reason.Code, payload.Reason.Message)
	return p.newAck(msg, payload.TaskID, "rejected")
}

// handleTaskNeedInfo handles information request from receiver
func (p *P2PProtocol) handleTaskNeedInfo(ctx context.Context, msg *types.UnifiedSynapseMessage) (*types.UnifiedSynapseMessage, error) {
	var payload types.TaskNeedInfoPayload
	if err := json.Unmarshal(msg.Payload, &payload); err != nil {
		return p.newNACK(msg, types.NewSynapseError(types.ErrMessageFormatInvalid, "Invalid need info payload", err.Error()))
	}

	log.Printf("[P2P] Information requested: %s (missing %d fields)", payload.TaskID, len(payload.MissingFields))

	// This would typically be handled by the application logic
	// For now, just acknowledge
	return p.newAck(msg, payload.TaskID, "info_request_received")
}

// handleTaskInfoProvided handles information provided by sender
func (p *P2PProtocol) handleTaskInfoProvided(ctx context.Context, msg *types.UnifiedSynapseMessage) (*types.UnifiedSynapseMessage, error) {
	var payload types.TaskInfoProvidedPayload
	if err := json.Unmarshal(msg.Payload, &payload); err != nil {
		return p.newNACK(msg, types.NewSynapseError(types.ErrMessageFormatInvalid, "Invalid info provided payload", err.Error()))
	}

	log.Printf("[P2P] Information provided: %s (%d fields)", payload.TaskID, len(payload.ProvidedInfo))
	return p.newAck(msg, payload.TaskID, "info_received")
}

// handleTaskStatusUpdate handles status updates from receiver
func (p *P2PProtocol) handleTaskStatusUpdate(ctx context.Context, msg *types.UnifiedSynapseMessage) (*types.UnifiedSynapseMessage, error) {
	var payload types.TaskStatusUpdatePayload
	if err := json.Unmarshal(msg.Payload, &payload); err != nil {
		return p.newNACK(msg, types.NewSynapseError(types.ErrMessageFormatInvalid, "Invalid status update payload", err.Error()))
	}

	log.Printf("[P2P] Status update: %s - %s (progress: %d%%)", payload.TaskID, payload.Status, payload.Progress)
	return p.newAck(msg, payload.TaskID, "status_received")
}

// handleTaskComplete handles task completion notification
func (p *P2PProtocol) handleTaskComplete(ctx context.Context, msg *types.UnifiedSynapseMessage) (*types.UnifiedSynapseMessage, error) {
	var payload types.TaskCompletePayload
	if err := json.Unmarshal(msg.Payload, &payload); err != nil {
		return p.newNACK(msg, types.NewSynapseError(types.ErrMessageFormatInvalid, "Invalid complete payload", err.Error()))
	}

	log.Printf("[P2P] Task completed: %s", payload.TaskID)
	return p.newAck(msg, payload.TaskID, "completed")
}

// handleTaskFailed handles task failure notification
func (p *P2PProtocol) handleTaskFailed(ctx context.Context, msg *types.UnifiedSynapseMessage) (*types.UnifiedSynapseMessage, error) {
	var payload types.TaskFailedPayload
	if err := json.Unmarshal(msg.Payload, &payload); err != nil {
		return p.newNACK(msg, types.NewSynapseError(types.ErrMessageFormatInvalid, "Invalid failed payload", err.Error()))
	}

	log.Printf("[P2P] Task failed: %s (%s: %s)", payload.TaskID, payload.Error.Code, payload.Error.Message)
	return p.newAck(msg, payload.TaskID, "failure_noted")
}

// handleTaskCancel handles task cancellation
func (p *P2PProtocol) handleTaskCancel(ctx context.Context, msg *types.UnifiedSynapseMessage) (*types.UnifiedSynapseMessage, error) {
	var payload types.TaskCancelPayload
	if err := json.Unmarshal(msg.Payload, &payload); err != nil {
		return p.newNACK(msg, types.NewSynapseError(types.ErrMessageFormatInvalid, "Invalid cancel payload", err.Error()))
	}

	log.Printf("[P2P] Task cancelled: %s (by %s)", payload.TaskID, payload.CancelledBy)
	return p.newAck(msg, payload.TaskID, "cancelled")
}

// handleTaskAck handles task-level acknowledgment
func (p *P2PProtocol) handleTaskAck(ctx context.Context, msg *types.UnifiedSynapseMessage) (*types.UnifiedSynapseMessage, error) {
	log.Printf("[P2P] Task ACK received for message: %s", msg.ReplyTo)
	return nil, nil // No response needed for ACK
}

// handleTaskNeedHuman handles human intervention request
func (p *P2PProtocol) handleTaskNeedHuman(ctx context.Context, msg *types.UnifiedSynapseMessage) (*types.UnifiedSynapseMessage, error) {
	var payload types.TaskNeedHumanPayload
	if err := json.Unmarshal(msg.Payload, &payload); err != nil {
		return p.newNACK(msg, types.NewSynapseError(types.ErrMessageFormatInvalid, "Invalid need human payload", err.Error()))
	}

	log.Printf("[P2P] Human intervention needed: %s (level: %d)", payload.TaskID, payload.EscalationLevel)

	// Store the HiL request and wait for human decision
	// This would typically be handled by a HiL manager

	return p.newAck(msg, payload.TaskID, "human_intervention_requested")
}

// handleHumanDecision handles human decision response
func (p *P2PProtocol) handleHumanDecision(ctx context.Context, msg *types.UnifiedSynapseMessage) (*types.UnifiedSynapseMessage, error) {
	var payload types.HumanDecisionPayload
	if err := json.Unmarshal(msg.Payload, &payload); err != nil {
		return p.newNACK(msg, types.NewSynapseError(types.ErrMessageFormatInvalid, "Invalid human decision payload", err.Error()))
	}

	log.Printf("[P2P] Human decision: %s (action: %s, by %s)", payload.TaskID, payload.Decision.Action, payload.DecidedBy.UserName)

	// Process the human decision
	// This would resume the task execution with the decision

	return p.newAck(msg, payload.TaskID, "decision_received")
}

// handleAck handles generic ACK
func (p *P2PProtocol) handleAck(ctx context.Context, msg *types.UnifiedSynapseMessage) (*types.UnifiedSynapseMessage, error) {
	log.Printf("[P2P] ACK received for: %s", msg.ReplyTo)
	return nil, nil
}

// handleNack handles generic NACK
func (p *P2PProtocol) handleNack(ctx context.Context, msg *types.UnifiedSynapseMessage) (*types.UnifiedSynapseMessage, error) {
	log.Printf("[P2P] NACK received for: %s", msg.ReplyTo)
	return nil, nil
}

// handleHeartbeat handles heartbeat
func (p *P2PProtocol) handleHeartbeat(ctx context.Context, msg *types.UnifiedSynapseMessage) (*types.UnifiedSynapseMessage, error) {
	return p.newAck(msg, "", "alive")
}

// =============================================================================
// Response Builders
// =============================================================================

// newAck creates an ACK response
func (p *P2PProtocol) newAck(msg *types.UnifiedSynapseMessage, taskID, note string) (*types.UnifiedSynapseMessage, error) {
	return &types.UnifiedSynapseMessage{
		ID:            types.GenerateMessageID(),
		Version:       "1.0.0",
		Type:          string(types.P2PAck),
		Sender:        p.agentEndpoint,
		Receiver:      &msg.Sender,
		CorrelationID: msg.CorrelationID,
		ReplyTo:       msg.ID,
		Timestamp:     types.NowTimestamp(),
	}, nil
}

// newNACK creates a NACK response
func (p *P2PProtocol) newNACK(msg *types.UnifiedSynapseMessage, err *types.SynapseError) (*types.UnifiedSynapseMessage, error) {
	return nil, err
}

// newReject creates a task rejection response
func (p *P2PProtocol) newReject(msg *types.UnifiedSynapseMessage, taskID, reasonCode, reasonMessage string, alternatives *struct {
	SuggestedAgentID string `json:"suggested_agent_id,omitempty"`
	RetryAfter       int    `json:"retry_after,omitempty"`
	PartialSupport   bool   `json:"partial_support,omitempty"`
}) (*types.UnifiedSynapseMessage, error) {
	payload := types.TaskRejectPayload{}
	payload.TaskID = taskID
	payload.Reason.Code = reasonCode
	payload.Reason.Message = reasonMessage
	if alternatives != nil {
		payload.Alternatives = alternatives
	}

	payloadBytes, _ := json.Marshal(payload)

	return &types.UnifiedSynapseMessage{
		ID:            types.GenerateMessageID(),
		Version:       "1.0.0",
		Type:          string(types.P2PTaskReject),
		Sender:        p.agentEndpoint,
		Receiver:      &msg.Sender,
		CorrelationID: msg.CorrelationID,
		ReplyTo:       msg.ID,
		Payload:       payloadBytes,
		Timestamp:     types.NowTimestamp(),
	}, nil
}

// newNeedInfo creates a need info response
func (p *P2PProtocol) newNeedInfo(msg *types.UnifiedSynapseMessage, taskID string, missingFields []types.MissingField) (*types.UnifiedSynapseMessage, error) {
	payload := types.TaskNeedInfoPayload{}
	payload.TaskID = taskID
	payload.MissingFields = missingFields

	payloadBytes, _ := json.Marshal(payload)

	return &types.UnifiedSynapseMessage{
		ID:            types.GenerateMessageID(),
		Version:       "1.0.0",
		Type:          string(types.P2PTaskNeedInfo),
		Sender:        p.agentEndpoint,
		Receiver:      &msg.Sender,
		CorrelationID: msg.CorrelationID,
		ReplyTo:       msg.ID,
		Payload:       payloadBytes,
		Timestamp:     types.NowTimestamp(),
	}, nil
}

// =============================================================================
// Validation
// =============================================================================

func (p *P2PProtocol) validateMessage(msg *types.UnifiedSynapseMessage) error {
	if msg.ID == "" {
		return types.ErrMessageFormatInvalidError
	}
	if msg.Sender.AgentID == "" {
		return types.NewSynapseError(types.ErrMessageFormatInvalid, "Missing sender agent_id", nil)
	}
	if msg.Sender.Endpoint == "" {
		return types.NewSynapseError(types.ErrMessageFormatInvalid, "Missing sender endpoint", nil)
	}
	return nil
}

// =============================================================================
// Task Management
// =============================================================================

// GetTaskStatus returns the status of a task
func (p *P2PProtocol) GetTaskStatus(taskID string) (types.P2PTaskStatus, bool) {
	p.mu.RLock()
	defer p.mu.RUnlock()
	status, ok := p.tasks[taskID]
	if !ok {
		return "", false
	}
	return *status, true
}

// ListTasks returns all tasks
func (p *P2PProtocol) ListTasks() map[string]types.P2PTaskStatus {
	p.mu.RLock()
	defer p.mu.RUnlock()
	result := make(map[string]types.P2PTaskStatus)
	for k, v := range p.tasks {
		result[k] = *v
	}
	return result
}