<?php

declare(strict_types=1);

namespace Drupal\flowdrop_stategraph\DTO;

use Drupal\flowdrop_orchestration\DTO\Snapshot\IteratorState;
use Drupal\flowdrop_orchestration\DTO\Snapshot\Message;
use Drupal\flowdrop_orchestration\DTO\Snapshot\WorkflowSnapshot;
use Drupal\flowdrop_stategraph\Attribute\StateField;

/**
 * Generic graph state container for stateful workflow execution.
 *
 * This class provides a flexible, generic state structure that can be used
 * for any stategraph workflow. It includes commonly needed fields for:
 * - Message history (for chat/agent workflows).
 * - Data accumulation (for processing pipelines).
 * - Iteration tracking (for loop control).
 * - Metadata storage (for debugging and audit).
 *
 * The state is designed to be serializable to JSON for checkpointing
 * and supports reducer-based updates for each field.
 *
 * Example usage:
 * @code
 * $state = new GraphState();
 * $state = $state->withUpdate(StateUpdate::create([
 *   "messages" => [new Message("user", "Hello")],
 *   "data" => ["key" => "value"],
 * ]));
 * @endcode
 *
 * @see \Drupal\flowdrop_stategraph\Attribute\StateField
 *
 * @api
 */
final class GraphState implements \JsonSerializable {

  /**
   * Constructs a new GraphState.
   *
   * @param array<Message> $messages
   *   Message history for chat/agent workflows.
   * @param array<string, mixed> $data
   *   General data storage for workflow state.
   * @param array<string, mixed> $metadata
   *   Execution metadata (timestamps, node history, etc.).
   * @param \Drupal\flowdrop_orchestration\DTO\Snapshot\IteratorState|null $iterator
   *   Iterator state for array processing loops.
   * @param string $currentNodeId
   *   The ID of the currently executing node.
   * @param int $iterationCount
   *   Current iteration count for loop protection.
   * @param string $threadId
   *   Thread ID for grouping related checkpoints.
   * @param bool $isComplete
   *   Whether the graph execution is complete.
   * @param string|null $error
   *   Error message if execution failed.
   */
  public function __construct(
    #[StateField(reducer: "append", description: "Message history for chat/agent workflows")]
    public readonly array $messages = [],
    #[StateField(reducer: "merge", description: "General data storage for workflow state")]
    public readonly array $data = [],
    #[StateField(reducer: "merge", description: "Execution metadata")]
    public readonly array $metadata = [],
    #[StateField(reducer: "replace", description: "Iterator state for array processing")]
    public readonly ?IteratorState $iterator = NULL,
    #[StateField(reducer: "replace", description: "Current executing node ID")]
    public readonly string $currentNodeId = "",
    #[StateField(reducer: "replace", description: "Current iteration count")]
    public readonly int $iterationCount = 0,
    #[StateField(reducer: "replace", description: "Thread ID for checkpoint grouping")]
    public readonly string $threadId = "",
    #[StateField(reducer: "replace", description: "Whether execution is complete")]
    public readonly bool $isComplete = FALSE,
    #[StateField(reducer: "replace", description: "Error message if failed")]
    public readonly ?string $error = NULL,
  ) {}

  /**
   * Creates a new GraphState with applied updates.
   *
   * This method applies a StateUpdate using the configured reducers
   * for each field, returning a new immutable state instance.
   *
   * @param \Drupal\flowdrop_stategraph\DTO\StateUpdate $update
   *   The state update to apply.
   * @param array<string, \Drupal\flowdrop_stategraph\Reducer\ReducerInterface> $reducers
   *   Map of reducer ID to reducer instance.
   *
   * @return self
   *   A new GraphState with the updates applied.
   */
  public function withUpdate(StateUpdate $update, array $reducers): self {
    $reflection = new \ReflectionClass($this);
    $newValues = [];

    foreach ($reflection->getProperties() as $property) {
      $propertyName = $property->getName();
      $currentValue = $this->{$propertyName};

      // Check if update has a value for this property.
      $updateData = $update->getData();
      if (!array_key_exists($propertyName, $updateData)) {
        $newValues[$propertyName] = $currentValue;
        continue;
      }

      // Get the StateField attribute to determine reducer.
      $attributes = $property->getAttributes(StateField::class);
      $reducerType = "replace";

      if (!empty($attributes)) {
        $stateField = $attributes[0]->newInstance();
        $reducerType = $stateField->reducer;
      }

      // Apply the reducer.
      $updateValue = $updateData[$propertyName];

      // A single message passed as a bare associative array would otherwise
      // be appended value-by-value by the append reducer; wrap it so it is
      // treated as one item.
      if ($propertyName === "messages" && is_array($updateValue) && isset($updateValue["role"])) {
        $updateValue = [$updateValue];
      }

      $reducer = $reducers[$reducerType] ?? $reducers["replace"];
      $newValues[$propertyName] = $reducer->reduce($currentValue, $updateValue);
    }

    return new self(
      // Re-normalize after reduction: processors may append plain arrays
      // (e.g. message data that round-tripped through output serialization),
      // and consumers rely on this channel containing Message objects only —
      // the same guarantee fromArray() makes.
      messages: self::normalizeMessages($newValues["messages"] ?? []),
      data: $newValues["data"] ?? [],
      metadata: $newValues["metadata"] ?? [],
      iterator: self::normalizeIterator($newValues["iterator"] ?? NULL),
      currentNodeId: $newValues["currentNodeId"] ?? "",
      iterationCount: $newValues["iterationCount"] ?? 0,
      threadId: $newValues["threadId"] ?? "",
      isComplete: $newValues["isComplete"] ?? FALSE,
      error: $newValues["error"] ?? NULL,
    );
  }

  /**
   * Creates a new state with incremented iteration count.
   *
   * @return self
   *   A new GraphState with iteration count incremented.
   */
  public function withIncrementedIteration(): self {
    return new self(
      messages: $this->messages,
      data: $this->data,
      metadata: $this->metadata,
      iterator: $this->iterator,
      currentNodeId: $this->currentNodeId,
      iterationCount: $this->iterationCount + 1,
      threadId: $this->threadId,
      isComplete: $this->isComplete,
      error: $this->error,
    );
  }

  /**
   * Creates a new state marked as complete.
   *
   * @return self
   *   A new GraphState marked as complete.
   */
  public function withComplete(): self {
    return new self(
      messages: $this->messages,
      data: $this->data,
      metadata: $this->metadata,
      iterator: $this->iterator,
      currentNodeId: $this->currentNodeId,
      iterationCount: $this->iterationCount,
      threadId: $this->threadId,
      isComplete: TRUE,
      error: $this->error,
    );
  }

  /**
   * Creates a new state with an error.
   *
   * @param string $errorMessage
   *   The error message.
   *
   * @return self
   *   A new GraphState with the error set.
   */
  public function withError(string $errorMessage): self {
    return new self(
      messages: $this->messages,
      data: $this->data,
      metadata: $this->metadata,
      iterator: $this->iterator,
      currentNodeId: $this->currentNodeId,
      iterationCount: $this->iterationCount,
      threadId: $this->threadId,
      isComplete: TRUE,
      error: $errorMessage,
    );
  }

  /**
   * Creates a GraphState from an array (e.g., from JSON).
   *
   * This method handles type normalization for:
   * - Messages: Converts arrays to Message objects.
   * - Iterator: Converts arrays to IteratorState objects.
   *
   * @param array<string, mixed> $data
   *   The state data array.
   *
   * @return self
   *   A new GraphState instance.
   */
  public static function fromArray(array $data): self {
    return new self(
      messages: self::normalizeMessages($data["messages"] ?? []),
      data: $data["data"] ?? [],
      metadata: $data["metadata"] ?? [],
      iterator: self::normalizeIterator($data["iterator"] ?? NULL),
      currentNodeId: (string) ($data["currentNodeId"] ?? ""),
      iterationCount: (int) ($data["iterationCount"] ?? 0),
      threadId: (string) ($data["threadId"] ?? ""),
      isComplete: (bool) ($data["isComplete"] ?? FALSE),
      error: isset($data["error"]) ? (string) $data["error"] : NULL,
    );
  }

  /**
   * Creates a GraphState from a WorkflowSnapshot.
   *
   * This factory method maps WorkflowSnapshot properties to GraphState
   * properties, handling type conversions:
   * - initialInput → data
   * - status (string) → isComplete (bool)
   * - iterator (array) → iterator (IteratorState)
   * - messages (array<array>) → messages (array<Message>)
   *
   * @param \Drupal\flowdrop_orchestration\DTO\Snapshot\WorkflowSnapshot $snapshot
   *   The WorkflowSnapshot to convert.
   * @param string|null $threadIdOverride
   *   Optional thread ID to use instead of the snapshot's value.
   *
   * @return self
   *   A new GraphState instance.
   */
  public static function fromWorkflowSnapshot(WorkflowSnapshot $snapshot, ?string $threadIdOverride = NULL): self {
    // Map WorkflowSnapshot status to GraphState isComplete.
    $isComplete = in_array($snapshot->status, [
      WorkflowSnapshot::STATUS_COMPLETED,
      WorkflowSnapshot::STATUS_FAILED,
      WorkflowSnapshot::STATUS_CANCELLED,
    ], TRUE);

    // Get currentNodeId from metadata if available.
    $currentNodeId = $snapshot->metadata["current_node_id"] ?? "";

    return new self(
      messages: self::normalizeMessages($snapshot->messages),
      data: $snapshot->initialInput,
      metadata: $snapshot->metadata,
      iterator: self::normalizeIterator($snapshot->iterator),
      currentNodeId: (string) $currentNodeId,
      iterationCount: $snapshot->iterationCount,
      threadId: $threadIdOverride ?? $snapshot->threadId,
      isComplete: $isComplete,
      error: $snapshot->error,
    );
  }

  /**
   * Converts GraphState data for WorkflowSnapshot compatibility.
   *
   * Returns an array with data formatted for WorkflowSnapshot construction.
   * This maps GraphState properties to WorkflowSnapshot naming conventions
   * and converts DTO objects to arrays.
   *
   * @return array<string, mixed>
   *   Data array compatible with WorkflowSnapshot construction.
   */
  public function toSnapshotData(): array {
    return [
      "messages" => self::serializeMessages($this->messages),
      "iterator" => $this->iterator?->toArray(),
      "iterationCount" => $this->iterationCount,
      "threadId" => $this->threadId,
      "error" => $this->error,
      "initialInput" => $this->data,
      "metadata" => array_merge($this->metadata, [
        "current_node_id" => $this->currentNodeId,
      ]),
    ];
  }

  /**
   * Normalizes messages from mixed input to Message objects.
   *
   * @param array<mixed> $messages
   *   Array of Message objects or arrays.
   *
   * @return array<Message>
   *   Array of Message objects.
   */
  public static function normalizeMessages(array $messages): array {
    $normalized = [];
    foreach ($messages as $messageData) {
      if ($messageData instanceof Message) {
        $normalized[] = $messageData;
      }
      elseif (is_array($messageData)) {
        $normalized[] = Message::fromArray($messageData);
      }
    }
    return $normalized;
  }

  /**
   * Serializes Message objects to arrays.
   *
   * @param array<Message> $messages
   *   Array of Message objects.
   *
   * @return array<array<string, mixed>>
   *   Array of message arrays.
   */
  public static function serializeMessages(array $messages): array {
    return array_map(
      static fn(Message $message): array => $message->toArray(),
      $messages
    );
  }

  /**
   * Normalizes iterator from mixed input to IteratorState or NULL.
   *
   * @param mixed $iterator
   *   IteratorState object, array, or NULL.
   *
   * @return \Drupal\flowdrop_orchestration\DTO\Snapshot\IteratorState|null
   *   IteratorState object or NULL.
   */
  public static function normalizeIterator(mixed $iterator): ?IteratorState {
    if ($iterator instanceof IteratorState) {
      return $iterator;
    }
    if (is_array($iterator)) {
      return IteratorState::fromArray($iterator);
    }
    return NULL;
  }

  /**
   * Converts the state to an array.
   *
   * @return array<string, mixed>
   *   The state as an associative array.
   */
  public function toArray(): array {
    return [
      "messages" => array_map(
        fn(Message $m) => $m->toArray(),
        $this->messages
      ),
      "data" => $this->data,
      "metadata" => $this->metadata,
      "iterator" => $this->iterator?->toArray(),
      "currentNodeId" => $this->currentNodeId,
      "iterationCount" => $this->iterationCount,
      "threadId" => $this->threadId,
      "isComplete" => $this->isComplete,
      "error" => $this->error,
    ];
  }

  /**
   * {@inheritdoc}
   */
  public function jsonSerialize(): mixed {
    return $this->toArray();
  }

}
