#pragma once

#include "llvm/Support/JSON.h"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp"
#include <cstdint>
#include <string>
#include <vector>

namespace onnx_mlir {
struct PimInstructionTiming;

enum class PimQueueTermination {
  NotRun, IntervalConverged, RecurrenceProved, EventBudget, GenerationLimit,
  HistoryLimit, NonMonotonic, NoReadyEvents
};

struct PimQueueEventLocation {
  int64_t core = 0;
  uint32_t pc = 0;
  std::string kind;
};

struct PimQueueProgress {
  size_t processedEvents = 0, completedGenerations = 0;
  uint64_t maximumQueueWaitNs = 0;
  std::optional<PimQueueEventLocation> maximumQueueRequest, generationLimit;
  std::vector<std::pair<int64_t, double>> perCoreQueueWaitNs;
  std::vector<uint64_t> completedIntervalsNs;
};

struct PimQueueIntervalWindow {
  std::optional<double> intervalNs;
  std::string reason;
  PimQueueTermination termination = PimQueueTermination::NotRun;
  PimQueueProgress progress;
};

struct PimQueuePolicyResult {
  PimQueueIntervalWindow legacy;
  spatial::ScheduleRecurrenceEvidence recurrence;
  PimQueueProgress progress;
  PimQueueTermination termination = PimQueueTermination::NotRun;
  size_t checkpoints = 0, historyBytes = 0, stateWords = 0, proofProcessedEvents = 0;
  bool reverseTies = false;
  std::optional<double> intervalNs() const;
};

struct PimQueueEvaluation {
  std::array<PimQueuePolicyResult, 2> policies;
  uint64_t memoryServiceNs = 0;
  double evaluationSeconds = 0;
  bool recurrenceRequired = false, details = false, measureTime = false;
  spatial::WholeScheduleCost cost() const;
};

llvm::json::Object serializePimQueueEvaluation(const PimQueueEvaluation &evaluation);

struct PimPeriodicEstimate {
  spatial::WholeScheduleCost cost;
  llvm::json::Object report;
  std::optional<PimQueueEvaluation> queue;
};

// Binary instructions, indexed by physical core. No executable or measured data.
struct PimInstructionProgram {
  std::vector<std::string> instructions;
  std::vector<std::vector<int64_t>> arrayGroups;
};

// Hardware work only. Bytes are already in the target timing precision.
// Repeat keeps ordered request patterns compact until queue construction.
struct PimExecutionEvent {
  enum Kind { Work, Load, Store, Send, Receive, Signal, Wait, Repeat };
  enum WorkClass { Matrix, Vector, LocalCopy, Fill, VectorCopy, ScalarIssue };
  Kind kind = Work;
  uint64_t amount = 0;
  int64_t peer = 0, reg = 0, count = 0;
  std::vector<PimExecutionEvent> body;
  bool constant = false;
  WorkClass workClass = Fill;
  PimExecutionEvent(Kind kind, uint64_t amount, int64_t peer = 0,
      int64_t reg = 0, int64_t count = 0, std::vector<PimExecutionEvent> body = {})
      : kind(kind), amount(amount), peer(peer), reg(reg), count(count), body(std::move(body)) {}
};
using PimExecutionProgram = std::vector<std::vector<PimExecutionEvent>>;

struct PimPeriodicOptions {
  bool networkQueues = true;
  bool instructionIssue = true;
  bool scalarIssue = true;
  bool ignoreLocalCopies = false;
  bool workBreakdown = false;
  bool measureTime = false;
  // Optional ordered work/IO extraction for timing validation; no execution.
  PimExecutionProgram *executionTrace = nullptr;
  // Prove recurrence of the modeled state. Instruction diagnostics retain
  // their legacy interval-window result separately.
  bool verifyPeriodicState = false;
  // Optional diagnostic cap. The complete-state proof keeps its original budget.
  std::optional<size_t> intervalWindowEventBudget;
  bool memoryController = true;
  bool detailed = false;
  // Apply the oracle's fp32-to-int8 timing projection without changing the program.
  bool oraclePrecision = false;
};

llvm::json::Object analyzePimExecution(const PimExecutionProgram &program,
    const PimInstructionTiming &target, PimPeriodicOptions options = {});
PimPeriodicEstimate evaluatePimExecution(const PimExecutionProgram &program,
    const PimInstructionTiming &target, PimPeriodicOptions options = {});

uint64_t getPimLocalTransferWorkNs(const PimInstructionTiming &target,
    uint64_t bytes, bool read, bool write, bool instructionIssue = true);
uint64_t getPimVectorWorkNs(const PimInstructionTiming &target,
    uint64_t elements, uint64_t inputBytes, uint64_t outputBytes,
    bool binary, bool instructionIssue = true);
uint64_t getPimMatrixOutputBytes(const PimInstructionTiming &target,
    uint64_t cellBits, uint64_t weightBits, uint64_t groups, uint64_t outputBytes);
uint64_t getPimMatrixWorkNs(const PimInstructionTiming &target,
    uint64_t inputBits, uint64_t outputBytes, bool instructionIssue = true);

llvm::json::Object analyzePimInstructions(const PimInstructionProgram &program,
    const PimInstructionTiming &target, int64_t cellBits,
    PimPeriodicOptions options = {});
} // namespace onnx_mlir