<?php

declare(strict_types=1);

namespace Drupal\flowdrop_runtime\Service\WorkflowState;

use Drupal\flowdrop_orchestration\DTO\Snapshot\NodeSnapshot;
use Drupal\flowdrop_runtime\DTO\Snapshot\ValidationResult;
use Drupal\flowdrop_orchestration\DTO\Snapshot\WorkflowSnapshot;
use Psr\Log\LoggerInterface;

/**
 * Default implementation of WorkflowStateManagerInterface.
 *
 * This implementation provides:
 * - In-memory caching of snapshots for fast access during execution.
 * - Persistent storage via EntitySnapshotStorage.
 * - Validation of snapshots against workflow definitions.
 * - Support for orchestrator handover with injected node states.
 * - Dependency tracking and validation.
 *
 * @internal This class is not part of the public API and may change without notice.
 *
 * @see \Drupal\flowdrop_runtime\Service\WorkflowState\WorkflowStateManagerInterface
 */
class WorkflowStateManager implements WorkflowStateManagerInterface {

  /**
   * In-memory cache of snapshots by execution ID.
   *
   * @var array<string, WorkflowSnapshot>
   */
  private array $snapshots = [];

  /**
   * Index of execution IDs by thread ID.
   *
   * @var array<string, array<string>>
   */
  private array $threadIndex = [];

  /**
   * Index of execution IDs by workflow ID.
   *
   * @var array<string, array<string>>
   */
  private array $workflowIndex = [];

  /**
   * Execution order counter for node completion tracking.
   *
   * @var array<string, int>
   */
  private array $executionOrderCounters = [];

  /**
   * Constructs a new WorkflowStateManager.
   *
   * @param \Psr\Log\LoggerInterface $logger
   *   The logger service.
   * @param \Drupal\flowdrop_runtime\Service\WorkflowState\EntitySnapshotStorage|null $entityStorage
   *   The entity storage for persistent snapshots.
   */
  public function __construct(
    private readonly LoggerInterface $logger,
    private readonly ?EntitySnapshotStorage $entityStorage = NULL,
  ) {}

  /**
   * {@inheritdoc}
   */
  public function createSnapshot(
    string $workflowId,
    string $workflowVersion,
    array $initialInput = [],
    ?string $threadId = NULL,
    array $metadata = [],
  ): WorkflowSnapshot {
    $snapshot = WorkflowSnapshot::create(
      workflowId: $workflowId,
      workflowVersion: $workflowVersion,
      initialInput: $initialInput,
      threadId: $threadId,
      metadata: $metadata,
    );

    $this->logger->debug("Created workflow snapshot @exec for workflow @workflow", [
      "@exec" => $snapshot->executionId,
      "@workflow" => $workflowId,
    ]);

    return $snapshot;
  }

  /**
   * {@inheritdoc}
   */
  public function initializeNodesFromWorkflow(
    WorkflowSnapshot $snapshot,
    array $workflowDefinition,
  ): WorkflowSnapshot {
    $nodes = $workflowDefinition["nodes"] ?? [];
    return $snapshot->withInitializedNodes($nodes);
  }

  /**
   * {@inheritdoc}
   */
  public function getSnapshot(string $executionId): ?WorkflowSnapshot {
    return $this->snapshots[$executionId] ?? NULL;
  }

  /**
   * {@inheritdoc}
   */
  public function getLatestByThread(string $threadId): ?WorkflowSnapshot {
    $executionIds = $this->threadIndex[$threadId] ?? [];
    if (empty($executionIds)) {
      return NULL;
    }

    // Get the most recent by checking updatedAt.
    $latest = NULL;
    $latestTime = 0;

    foreach ($executionIds as $executionId) {
      $snapshot = $this->snapshots[$executionId] ?? NULL;
      if ($snapshot !== NULL && $snapshot->updatedAt > $latestTime) {
        $latest = $snapshot;
        $latestTime = $snapshot->updatedAt;
      }
    }

    return $latest;
  }

  /**
   * {@inheritdoc}
   */
  public function getNodeState(string $executionId, string $nodeId): ?NodeSnapshot {
    $snapshot = $this->getSnapshot($executionId);
    return $snapshot?->getNodeState($nodeId);
  }

  /**
   * {@inheritdoc}
   */
  public function save(WorkflowSnapshot $snapshot): WorkflowSnapshot {
    $executionId = $snapshot->executionId;

    // Update the snapshot with current timestamp.
    $now = time();
    $updatedSnapshot = new WorkflowSnapshot(
      workflowId: $snapshot->workflowId,
      workflowVersion: $snapshot->workflowVersion,
      executionId: $snapshot->executionId,
      status: $snapshot->status,
      nodeStates: $snapshot->nodeStates,
      initialInput: $snapshot->initialInput,
      messages: $snapshot->messages,
      iterator: $snapshot->iterator,
      iterationCount: $snapshot->iterationCount,
      threadId: $snapshot->threadId,
      error: $snapshot->error,
      metadata: $snapshot->metadata,
      createdAt: $snapshot->createdAt ?: $now,
      updatedAt: $now,
    );

    // Store in memory cache.
    $this->snapshots[$executionId] = $updatedSnapshot;

    // Update indices.
    $this->updateIndices($updatedSnapshot);

    // Persist to entity storage if available.
    if ($this->entityStorage !== NULL) {
      try {
        $this->entityStorage->save($updatedSnapshot);
        $this->logger->debug("Persisted workflow snapshot @exec to entity storage", [
          "@exec" => $executionId,
        ]);
      }
      catch (\Exception $e) {
        $this->logger->error("Failed to persist workflow snapshot @exec: @error", [
          "@exec" => $executionId,
          "@error" => $e->getMessage(),
        ]);
      }
    }

    $this->logger->debug("Saved workflow snapshot @exec (status: @status)", [
      "@exec" => $executionId,
      "@status" => $updatedSnapshot->status,
    ]);

    return $updatedSnapshot;
  }

  /**
   * {@inheritdoc}
   */
  public function delete(string $executionId): bool {
    if (!isset($this->snapshots[$executionId])) {
      return FALSE;
    }

    $snapshot = $this->snapshots[$executionId];

    // Remove from indices.
    $this->removeFromIndices($snapshot);

    // Remove from storage.
    unset($this->snapshots[$executionId]);
    unset($this->executionOrderCounters[$executionId]);

    $this->logger->debug("Deleted workflow snapshot @exec", [
      "@exec" => $executionId,
    ]);

    return TRUE;
  }

  /**
   * {@inheritdoc}
   */
  public function updateStatus(string $executionId, string $status): WorkflowSnapshot {
    $snapshot = $this->getOrFail($executionId);
    $updated = $snapshot->withStatus($status);
    return $this->save($updated);
  }

  /**
   * {@inheritdoc}
   */
  public function setNodeRunning(string $executionId, string $nodeId): WorkflowSnapshot {
    $snapshot = $this->getOrFail($executionId);
    $order = $this->getNextExecutionOrder($executionId);
    $updated = $snapshot->withNodeRunning($nodeId, $order);
    return $this->save($updated);
  }

  /**
   * {@inheritdoc}
   */
  public function setNodeCompleted(
    string $executionId,
    string $nodeId,
    array $output,
    bool $injected = FALSE,
  ): WorkflowSnapshot {
    $snapshot = $this->getOrFail($executionId);
    $order = $this->getNextExecutionOrder($executionId);
    $updated = $snapshot->withNodeCompleted($nodeId, $output, $order, $injected);
    return $this->save($updated);
  }

  /**
   * {@inheritdoc}
   */
  public function setNodeFailed(
    string $executionId,
    string $nodeId,
    string $error,
  ): WorkflowSnapshot {
    $snapshot = $this->getOrFail($executionId);
    $order = $this->getNextExecutionOrder($executionId);
    $updated = $snapshot->withNodeFailed($nodeId, $error, $order);
    return $this->save($updated);
  }

  /**
   * {@inheritdoc}
   */
  public function setNodeSkipped(
    string $executionId,
    string $nodeId,
    string $reason = "",
  ): WorkflowSnapshot {
    $snapshot = $this->getOrFail($executionId);
    $updated = $snapshot->withNodeSkipped($nodeId, $reason);
    return $this->save($updated);
  }

  /**
   * {@inheritdoc}
   */
  public function markNodesCompleted(
    string $executionId,
    array $nodeOutputs,
  ): WorkflowSnapshot {
    $snapshot = $this->getOrFail($executionId);
    $startingOrder = $this->getNextExecutionOrder($executionId);

    // Reserve execution order slots.
    $this->executionOrderCounters[$executionId] = $startingOrder + count($nodeOutputs);

    $updated = $snapshot->withInjectedNodes($nodeOutputs, $startingOrder);
    return $this->save($updated);
  }

  /**
   * {@inheritdoc}
   */
  public function initializeFromSnapshot(
    WorkflowSnapshot $source,
    bool $preserveCompletedNodes = TRUE,
  ): WorkflowSnapshot {
    $newSnapshot = WorkflowSnapshot::createFromSnapshot($source, $preserveCompletedNodes);

    // Mark preserved nodes as injected.
    if ($preserveCompletedNodes) {
      $nodeStates = [];
      foreach ($newSnapshot->nodeStates as $nodeId => $nodeSnapshot) {
        if ($nodeSnapshot->isCompleted() && !$nodeSnapshot->injected) {
          // Mark as injected since it wasn't executed in this context.
          $nodeStates[$nodeId] = NodeSnapshot::injected(
            nodeId: $nodeId,
            output: $nodeSnapshot->output ?? [],
            executionOrder: $nodeSnapshot->executionOrder,
            metadata: array_merge($nodeSnapshot->metadata, [
              "source_execution_id" => $source->executionId,
            ]),
          );
        }
        else {
          $nodeStates[$nodeId] = $nodeSnapshot;
        }
      }

      $newSnapshot = new WorkflowSnapshot(
        workflowId: $newSnapshot->workflowId,
        workflowVersion: $newSnapshot->workflowVersion,
        executionId: $newSnapshot->executionId,
        status: $newSnapshot->status,
        nodeStates: $nodeStates,
        initialInput: $newSnapshot->initialInput,
        messages: $newSnapshot->messages,
        iterator: $newSnapshot->iterator,
        iterationCount: $newSnapshot->iterationCount,
        threadId: $newSnapshot->threadId,
        error: $newSnapshot->error,
        metadata: $newSnapshot->metadata,
        createdAt: $newSnapshot->createdAt,
        updatedAt: $newSnapshot->updatedAt,
      );
    }

    $this->logger->info("Initialized new snapshot @new from source @source", [
      "@new" => $newSnapshot->executionId,
      "@source" => $source->executionId,
    ]);

    return $newSnapshot;
  }

  /**
   * {@inheritdoc}
   */
  public function validateSnapshot(
    WorkflowSnapshot $snapshot,
    array $workflowDefinition,
  ): ValidationResult {
    $errors = [];
    $warnings = [];

    // 1. Validate workflow version.
    $expectedVersion = $this->computeWorkflowVersion($workflowDefinition);
    if ($snapshot->workflowVersion !== "" && $snapshot->workflowVersion !== $expectedVersion) {
      $errors[] = ValidationResult::error(
        code: ValidationResult::ERROR_VERSION_MISMATCH,
        message: "Workflow version mismatch. Expected: {$expectedVersion}, got: {$snapshot->workflowVersion}",
        details: [
          "expected" => $expectedVersion,
          "actual" => $snapshot->workflowVersion,
        ],
      );
    }

    // 2. Validate all node IDs exist in workflow.
    $workflowNodeIds = $this->extractNodeIds($workflowDefinition);
    foreach ($snapshot->nodeStates as $nodeId => $nodeSnapshot) {
      if (!in_array($nodeId, $workflowNodeIds, TRUE)) {
        $errors[] = ValidationResult::error(
          code: ValidationResult::ERROR_UNKNOWN_NODE,
          message: "Node '{$nodeId}' in snapshot does not exist in workflow definition",
          nodeId: $nodeId,
        );
      }
    }

    // 3. Check for orphan nodes in workflow not in snapshot (warning only).
    foreach ($workflowNodeIds as $nodeId) {
      if (!isset($snapshot->nodeStates[$nodeId])) {
        $warnings[] = ValidationResult::warning(
          code: ValidationResult::WARNING_ORPHAN_NODE,
          message: "Node '{$nodeId}' in workflow has no state in snapshot",
          nodeId: $nodeId,
        );
      }
    }

    // 4. Validate dependencies.
    $depResult = $this->validateDependencies($snapshot, $workflowDefinition);
    if (!$depResult->isValid()) {
      $errors = array_merge($errors, $depResult->getErrors());
    }
    $warnings = array_merge($warnings, $depResult->getWarnings());

    if (!empty($errors)) {
      return ValidationResult::invalid($errors, $warnings, [
        "validated_at" => time(),
        "workflow_id" => $snapshot->workflowId,
      ]);
    }

    return ValidationResult::valid($warnings, [
      "validated_at" => time(),
      "workflow_id" => $snapshot->workflowId,
    ]);
  }

  /**
   * {@inheritdoc}
   */
  public function validateDependencies(
    WorkflowSnapshot $snapshot,
    array $workflowDefinition,
  ): ValidationResult {
    $errors = [];
    $warnings = [];

    // Build dependency map from edges.
    $edges = $workflowDefinition["edges"] ?? [];
    $dependencies = $this->buildDependencyMap($edges);

    // For each completed node, check that all its dependencies are completed.
    foreach ($snapshot->nodeStates as $nodeId => $nodeSnapshot) {
      if (!$nodeSnapshot->isCompleted()) {
        continue;
      }

      $deps = $dependencies[$nodeId] ?? [];
      foreach ($deps as $depNodeId) {
        $depSnapshot = $snapshot->getNodeState($depNodeId);

        if ($depSnapshot === NULL) {
          $errors[] = ValidationResult::error(
            code: ValidationResult::ERROR_DEPENDENCY_NOT_MET,
            message: "Node '{$nodeId}' is completed but dependency '{$depNodeId}' has no state",
            nodeId: $nodeId,
            details: ["missing_dependency" => $depNodeId],
          );
        }
        elseif (!$depSnapshot->isCompleted() && !$depSnapshot->status === NodeSnapshot::STATUS_SKIPPED) {
          $errors[] = ValidationResult::error(
            code: ValidationResult::ERROR_DEPENDENCY_NOT_MET,
            message: "Node '{$nodeId}' is completed but dependency '{$depNodeId}' is {$depSnapshot->status}",
            nodeId: $nodeId,
            details: [
              "dependency_node" => $depNodeId,
              "dependency_status" => $depSnapshot->status,
            ],
          );
        }
      }

      // Warning if injected without verifying dependencies.
      if ($nodeSnapshot->injected && !empty($deps)) {
        $warnings[] = ValidationResult::warning(
          code: ValidationResult::WARNING_INJECTED_WITHOUT_DEPENDENCY,
          message: "Injected node '{$nodeId}' has dependencies that should be verified",
          nodeId: $nodeId,
          details: ["dependencies" => $deps],
        );
      }
    }

    if (!empty($errors)) {
      return ValidationResult::invalid($errors, $warnings);
    }

    return ValidationResult::valid($warnings);
  }

  /**
   * {@inheritdoc}
   */
  public function computeWorkflowVersion(array $workflowDefinition): string {
    // Create a normalized representation for hashing.
    $nodes = $workflowDefinition["nodes"] ?? [];
    $edges = $workflowDefinition["edges"] ?? [];

    // Sort nodes and edges by ID for consistent hashing.
    usort($nodes, fn($a, $b) => ($a["id"] ?? "") <=> ($b["id"] ?? ""));
    usort($edges, fn($a, $b) => ($a["id"] ?? "") <=> ($b["id"] ?? ""));

    // Extract only structural elements (not labels/positions).
    $structural = [
      "nodes" => array_map(fn($n) => [
        "id" => $n["id"] ?? "",
        "type" => $n["data"]["type"] ?? "",
      ], $nodes),
      "edges" => array_map(fn($e) => [
        "id" => $e["id"] ?? "",
        "source" => $e["source"] ?? "",
        "target" => $e["target"] ?? "",
        "sourceHandle" => $e["sourceHandle"] ?? "",
        "targetHandle" => $e["targetHandle"] ?? "",
      ], $edges),
    ];

    return substr(hash("sha256", json_encode($structural) ?: ""), 0, 16);
  }

  /**
   * {@inheritdoc}
   */
  public function addMessage(
    string $executionId,
    string $role,
    string $content,
    array $metadata = [],
  ): WorkflowSnapshot {
    $snapshot = $this->getOrFail($executionId);

    $message = [
      "role" => $role,
      "content" => $content,
      "timestamp" => time(),
    ];

    if (!empty($metadata)) {
      $message["metadata"] = $metadata;
    }

    $updated = $snapshot->withMessage($message);
    return $this->save($updated);
  }

  /**
   * {@inheritdoc}
   */
  public function updateIterator(
    string $executionId,
    ?array $iterator,
  ): WorkflowSnapshot {
    $snapshot = $this->getOrFail($executionId);
    $updated = $snapshot->withIterator($iterator);
    return $this->save($updated);
  }

  /**
   * {@inheritdoc}
   */
  public function incrementIteration(string $executionId): WorkflowSnapshot {
    $snapshot = $this->getOrFail($executionId);
    $updated = $snapshot->withIncrementedIteration();
    return $this->save($updated);
  }

  /**
   * {@inheritdoc}
   */
  public function listByWorkflow(
    string $workflowId,
    int $limit = 100,
    int $offset = 0,
  ): array {
    $executionIds = $this->workflowIndex[$workflowId] ?? [];

    // Get snapshots and sort by creation time (newest first).
    $snapshots = [];
    foreach ($executionIds as $executionId) {
      if (isset($this->snapshots[$executionId])) {
        $snapshots[] = $this->snapshots[$executionId];
      }
    }

    usort($snapshots, fn($a, $b) => $b->createdAt <=> $a->createdAt);

    return array_slice($snapshots, $offset, $limit);
  }

  /**
   * {@inheritdoc}
   */
  public function listByThread(string $threadId, int $limit = 100): array {
    $executionIds = $this->threadIndex[$threadId] ?? [];

    // Get snapshots and sort by creation time (newest first).
    $snapshots = [];
    foreach ($executionIds as $executionId) {
      if (isset($this->snapshots[$executionId])) {
        $snapshots[] = $this->snapshots[$executionId];
      }
    }

    usort($snapshots, fn($a, $b) => $b->createdAt <=> $a->createdAt);

    return array_slice($snapshots, 0, $limit);
  }

  /**
   * {@inheritdoc}
   */
  public function cleanup(int $olderThan, array $statuses = []): int {
    $deleted = 0;

    foreach ($this->snapshots as $executionId => $snapshot) {
      if ($snapshot->createdAt >= $olderThan) {
        continue;
      }

      if (!empty($statuses) && !in_array($snapshot->status, $statuses, TRUE)) {
        continue;
      }

      $this->delete($executionId);
      $deleted++;
    }

    $this->logger->info("Cleaned up @count workflow snapshots older than @time", [
      "@count" => $deleted,
      "@time" => date("Y-m-d H:i:s", $olderThan),
    ]);

    return $deleted;
  }

  /**
   * Gets a snapshot or throws an exception if not found.
   *
   * @param string $executionId
   *   The execution ID.
   *
   * @return \Drupal\flowdrop_orchestration\DTO\Snapshot\WorkflowSnapshot
   *   The snapshot.
   *
   * @throws \InvalidArgumentException
   *   If the snapshot is not found.
   */
  private function getOrFail(string $executionId): WorkflowSnapshot {
    $snapshot = $this->getSnapshot($executionId);
    if ($snapshot === NULL) {
      throw new \InvalidArgumentException(
        "Workflow snapshot not found: {$executionId}"
      );
    }
    return $snapshot;
  }

  /**
   * Gets the next execution order for a snapshot.
   *
   * @param string $executionId
   *   The execution ID.
   *
   * @return int
   *   The next execution order number.
   */
  private function getNextExecutionOrder(string $executionId): int {
    if (!isset($this->executionOrderCounters[$executionId])) {
      // Initialize from snapshot's max execution order.
      $snapshot = $this->getSnapshot($executionId);
      $this->executionOrderCounters[$executionId] = $snapshot !== NULL
        ? $snapshot->getMaxExecutionOrder() + 1
        : 1;
    }

    return $this->executionOrderCounters[$executionId]++;
  }

  /**
   * Updates the lookup indices for a snapshot.
   *
   * @param \Drupal\flowdrop_orchestration\DTO\Snapshot\WorkflowSnapshot $snapshot
   *   The snapshot to index.
   */
  private function updateIndices(WorkflowSnapshot $snapshot): void {
    $executionId = $snapshot->executionId;

    // Update thread index.
    if ($snapshot->threadId !== "") {
      if (!isset($this->threadIndex[$snapshot->threadId])) {
        $this->threadIndex[$snapshot->threadId] = [];
      }
      if (!in_array($executionId, $this->threadIndex[$snapshot->threadId], TRUE)) {
        $this->threadIndex[$snapshot->threadId][] = $executionId;
      }
    }

    // Update workflow index.
    if ($snapshot->workflowId !== "") {
      if (!isset($this->workflowIndex[$snapshot->workflowId])) {
        $this->workflowIndex[$snapshot->workflowId] = [];
      }
      if (!in_array($executionId, $this->workflowIndex[$snapshot->workflowId], TRUE)) {
        $this->workflowIndex[$snapshot->workflowId][] = $executionId;
      }
    }
  }

  /**
   * Removes a snapshot from the lookup indices.
   *
   * @param \Drupal\flowdrop_orchestration\DTO\Snapshot\WorkflowSnapshot $snapshot
   *   The snapshot to remove from indices.
   */
  private function removeFromIndices(WorkflowSnapshot $snapshot): void {
    $executionId = $snapshot->executionId;

    // Remove from thread index.
    if ($snapshot->threadId !== "" && isset($this->threadIndex[$snapshot->threadId])) {
      $this->threadIndex[$snapshot->threadId] = array_filter(
        $this->threadIndex[$snapshot->threadId],
        fn($id) => $id !== $executionId
      );
    }

    // Remove from workflow index.
    if ($snapshot->workflowId !== "" && isset($this->workflowIndex[$snapshot->workflowId])) {
      $this->workflowIndex[$snapshot->workflowId] = array_filter(
        $this->workflowIndex[$snapshot->workflowId],
        fn($id) => $id !== $executionId
      );
    }
  }

  /**
   * Extracts node IDs from a workflow definition.
   *
   * @param array<string, mixed> $workflowDefinition
   *   The workflow definition.
   *
   * @return array<string>
   *   Array of node IDs.
   */
  private function extractNodeIds(array $workflowDefinition): array {
    $nodes = $workflowDefinition["nodes"] ?? [];
    return array_map(fn($n) => $n["id"] ?? "", $nodes);
  }

  /**
   * Builds a dependency map from workflow edges.
   *
   * @param array<array<string, mixed>> $edges
   *   The workflow edges.
   *
   * @return array<string, array<string>>
   *   Map of node ID to array of dependency node IDs.
   */
  private function buildDependencyMap(array $edges): array {
    $dependencies = [];

    foreach ($edges as $edge) {
      $source = $edge["source"] ?? "";
      $target = $edge["target"] ?? "";

      if ($source === "" || $target === "") {
        continue;
      }

      if (!isset($dependencies[$target])) {
        $dependencies[$target] = [];
      }

      if (!in_array($source, $dependencies[$target], TRUE)) {
        $dependencies[$target][] = $source;
      }
    }

    return $dependencies;
  }

  /**
   * {@inheritdoc}
   */
  public function getSnapshotEntity(string $executionId): ?object {
    if ($this->entityStorage === NULL) {
      return NULL;
    }

    return $this->entityStorage->loadByExecutionId($executionId);
  }

}
