#include "PimPipelineSelection.hpp"
#include "PimCompilerOptions.hpp"
#include "PimCompilerUtils.hpp"
#include "PimSpatialCost.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp"
#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
#include "src/Accelerators/PIM/Common/Support/FileSystemUtils.hpp"
#include "src/Accelerators/PIM/Common/Support/ReportUtils.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
#include "mlir/Pass/PassManager.h"
#include "mlir/Pass/PassInstrumentation.h"
#include "src/Compiler/CompilerOptions.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.hpp"
#include "llvm/Support/FormatVariadic.h"
#include "llvm/Support/JSON.h"
#include "llvm/Support/raw_os_ostream.h"
#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
#include <chrono>

using namespace mlir;
using namespace onnx_mlir;
using namespace onnx_mlir::spatial;

namespace {
using Clock = std::chrono::steady_clock;
double elapsed(Clock::time_point start) {
  return std::chrono::duration<double>(Clock::now() - start).count();
}

llvm::json::Object countIR(Operation *root) {
  int64_t operations = 0, values = 0;
  root->walk([&](Operation *op) {
    ++operations;
    values += op->getNumResults();
    for (Region &region : op->getRegions())
      for (Block &block : region) values += block.getNumArguments();
  });
  return llvm::json::Object{{"operations", operations}, {"values", values}};
}

// Enabled with the existing compiler timing option. No instruction expansion.
struct CandidateTiming final : PassInstrumentation {
  llvm::json::Object &summary;
  Clock::time_point start;
  explicit CandidateTiming(llvm::json::Object &summary) : summary(summary) {}
  void runBeforePass(Pass *pass, Operation *op) override {
    if (pass->getArgument() == "pim-one-shot-bufferization")
      summary["bufferization_input_ir"] = countIR(op);
    start = Clock::now();
  }
  void runAfterPass(Pass *pass, Operation *) override {
    auto *phases = summary.getObject("lowering_phases_seconds");
    if (!phases) {
      summary["lowering_phases_seconds"] = llvm::json::Object{};
      phases = summary.getObject("lowering_phases_seconds");
    }
    (*phases)[pass->getName()] = elapsed(start);
  }
};

void writeSelectionReport(llvm::json::Object report) {
  auto file = openReportFileWithExtension("pipeline_selection", "json");
  if (file) {
    llvm::raw_os_ostream out(file);
    out << llvm::formatv("{0:2}\n", llvm::json::Value(std::move(report)));
  }
}

struct PimPipelineSelectionPass final : PassWrapper<PimPipelineSelectionPass, OperationPass<ModuleOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimPipelineSelectionPass)
  PimPipelineSelectionPass(const SchedulingTarget &target, const SpatialTargetResources &resources,
      size_t stages, SpatialDataflowExportStage exportStage)
      : target(target), resources(resources), stages(stages), exportStage(exportStage) {}
  SchedulingTarget target;
  SpatialTargetResources resources;
  size_t stages;
  SpatialDataflowExportStage exportStage;

  StringRef getArgument() const override { return "pim-pipeline-selection"; }
  StringRef getDescription() const override { return "Select a throughput pipeline using static Spatial queue analysis"; }

  LogicalResult realizeCandidate(ModuleOp checkpoint, const MergeScheduleResult &schedule,
      OwningOpRef<ModuleOp> &candidate, std::string &scheduleReport,
      llvm::json::Object &summary) {
    auto begin = Clock::now();
    IRMapping mapping;
    candidate = cast<ModuleOp>(checkpoint->clone(mapping));
    auto remapped = remapSpatialSchedule(schedule, mapping);
    if (failed(realizeSpatialSchedule(*candidate, remapped, target, stages,
          SpatialDataflowExportStage::None, nullptr, &scheduleReport))) return failure();
    summary["realization_seconds"] = elapsed(begin);
    return success();
  }

  WholeScheduleCost spatialCost(ModuleOp candidate, llvm::json::Object &summary) {
    WholeScheduleCost cost;
    auto begin = Clock::now();
    llvm::json::Object estimate;
    if (target.estimateRealizedScheduleCost)
      cost = target.estimateRealizedScheduleCost(candidate, estimate);
    else cost.reason = target.wholeScheduleCostUnavailable.empty()
      ? "target has no realized-Spatial cost query" : target.wholeScheduleCostUnavailable;
    cost.selectionMargin = kPipelineSelectionMargin;
    estimate["safe_for_selection"] = cost.eligible();
    estimate["diagnostic_only"] = false;
    summary["spatial_cost"] = std::move(estimate);
    summary["spatial_cost_seconds"] = elapsed(begin);
    return cost;
  }

  LogicalResult evaluateSpatialCandidates(ModuleOp checkpoint,
      PreparedSpatialSchedules &prepared, std::vector<WholeScheduleCost> &costs,
      llvm::json::Object &report) {
    llvm::json::Array summaries;
    double realization = 0, extraction = 0, expansion = 0, queue = 0, evaluation = 0, dependency = 0;
    auto evaluate = [&](const MergeScheduleResult &schedule) {
      size_t index = summaries.size();
      llvm::json::Object summary{{"index", static_cast<int64_t>(index)}};
      OwningOpRef<ModuleOp> candidate;
      WholeScheduleCost cost;
      std::string scheduleReport;
      {
        ScopedCompilerOutputSuppression silence;
        if (failed(realizeCandidate(checkpoint, schedule, candidate, scheduleReport, summary)))
          cost.reason = "Spatial schedule realization failed verification";
        else cost = spatialCost(*candidate, summary);
      }
      realization += summary.getNumber("realization_seconds").value_or(0);
      evaluation += summary.getNumber("spatial_cost_seconds").value_or(0);
      extraction += cost.timing.extractionSeconds;
      expansion += cost.timing.expansionSeconds;
      queue += cost.timing.evaluationSeconds;
      dependency += cost.timing.dependencySeconds;
      auto convergence = cost.convergenceProved(), support = cost.loweringSupported();
      summary["cost_available"] = cost.available();
      summary["convergence_proved"] = convergence
        ? llvm::json::Value(*convergence) : llvm::json::Value(nullptr);
      summary["lowering_supported"] = support
        ? llvm::json::Value(*support) : llvm::json::Value(nullptr);
      summary["eligible"] = cost.eligible();
      if (cost.primary) summary["interval_ns"] = *cost.primary;
      if (cost.alternative) summary["alternative_tie_interval_ns"] = *cost.alternative;
      llvm::json::Array budgets, witness;
      std::vector<size_t> cores(stages);
      std::vector<double> pressure(stages);
      for (size_t core = 0; core < schedule.processorStages.size(); ++core) {
        size_t stage = schedule.processorStages[core];
        ++cores[stage];
        if (core < cost.processorCriticalDelayNs.size())
          pressure[stage] += cost.processorCriticalDelayNs[core];
      }
      for (size_t stage = 0; stage < stages; ++stage) {
        budgets.push_back(static_cast<int64_t>(cores[stage]));
        witness.push_back(pressure[stage]);
      }
      summary["stage_core_budgets"] = std::move(budgets);
      summary["recurrence_stage_contributions_ns"] = std::move(witness);
      llvm::json::Array runs;
      for (size_t begin = 0; begin < prepared.graph.nodes.size();) {
        auto stageFor = [&](size_t task) {
          return schedule.processorStages[schedule.computeToCpuMap.lookup(
            prepared.graph.nodes[task].instance)];
        };
        size_t end = begin + 1, stage = stageFor(begin);
        while (end < prepared.graph.nodes.size() && stageFor(end) == stage) ++end;
        runs.push_back(llvm::json::Array{static_cast<int64_t>(begin),
          static_cast<int64_t>(end - begin), static_cast<int64_t>(stage)});
        begin = end;
      }
      summary["task_stage_runs"] = std::move(runs);
      if (!cost.eligible()) summary["reason"] = cost.reason;
      summaries.push_back(std::move(summary));
      return cost;
    };
    std::string error;
    auto start = Clock::now();
    auto search = optimizePipelineSchedule(prepared.graph, prepared.candidates.front(),
      stages, target, evaluate, error);
    report["search_seconds"] = elapsed(start);
    report["search_refinement_seconds"] = std::max(0.0, elapsed(start) - realization - evaluation);
    report["cost_evaluation_seconds"] = evaluation;
    report["candidate_realization_seconds"] = realization;
    report["event_extraction_seconds"] = extraction;
    report["event_expansion_seconds"] = expansion;
    report["queue_seconds"] = queue;
    report["dependency_analysis_seconds"] = dependency;
    report["candidates"] = std::move(summaries);
    if (failed(search)) {
      report["status"] = "incomplete-costs";
      report["reason"] = error;
      checkpoint.emitError("Spatial pipeline cannot compare the complete candidate set: ") << error;
      return failure();
    }
    prepared.candidates = std::move(search->schedules);
    costs = std::move(search->costs);
    bool single = costs.empty();
    if (single) {
      report["candidates"] = llvm::json::Array{llvm::json::Object{{"index", 0}}};
    } else {
      auto *entries = report.getArray("candidates");
      for (size_t i = 0; i < entries->size(); ++i)
        (*entries)[i].getAsObject()->insert({"parent", static_cast<int64_t>(search->parents[i])});
    }
    report["generated_candidates"] = static_cast<int64_t>(prepared.candidates.size());
    report["single_candidate_analysis_skipped"] = single;
    report["search_rounds"] = static_cast<int64_t>(search->rounds);
    report["packing_attempts"] = static_cast<int64_t>(search->packingAttempts);
    report["neighborhood_seconds"] = search->neighborhoodSeconds;
    report["constructive_packing_seconds"] = search->packingSeconds;
    report["considered_assignments"] = static_cast<int64_t>(search->consideredAssignments);
    report["search_termination"] = search->termination;
    report["queue_evaluation_budget"] = static_cast<int64_t>(kPipelineQueueEvaluationBudget);
    report["beam_width"] = static_cast<int64_t>(kPipelineBeamWidth);
    return success();
  }

  void runOnOperation() override {
    ModuleOp checkpoint = getOperation();
    auto start = Clock::now();
    auto prepared = prepareSpatialSchedules(checkpoint, target, stages);
    if (failed(prepared)) { signalPassFailure(); return; }
    if (prepared->candidates.empty()) {
      checkpoint.emitError("Spatial pipeline requires at least one candidate schedule");
      signalPassFailure(); return;
    }
    llvm::json::Object report{{"strategy", "spatial-queue-aware"}, {"static_analysis", true},
      {"candidate_generation_seconds", elapsed(start)},
      {"generated_candidates", static_cast<int64_t>(prepared->candidates.size())},
      {"selection_margin", kPipelineSelectionMargin},
      {"prediction_error_bound_available", false}, {"temporary_disk_bytes", 0}};
    if (pimDisableSynchronization) {
      report["synchronization_ablation"] = true;
      report["cost_model_program"] = "synchronized-reference";
      llvm::errs() << "Synchronization ablation: keep the normal selected schedule; "
        "report zero-cost SYNC/WAIT separately and omit them in code generation. "
        "Buffer reuse and restart correctness are not guaranteed.\n";
    }
    std::vector<WholeScheduleCost> costs;
    if (failed(evaluateSpatialCandidates(checkpoint, *prepared, costs, report))) {
      writeSelectionReport(std::move(report)); signalPassFailure(); return;
    }
    bool single = costs.empty();
    OwningOpRef<ModuleOp> winner;
    std::string scheduleReport;
    llvm::json::Array attempts;
    double realizationSeconds = 0, backendSeconds = 0;
    size_t lowered = 0;
    auto attempt = [&](size_t index) {
      llvm::json::Object summary{{"candidate", static_cast<int64_t>(index)}};
      std::string capturedReport;
      OwningOpRef<ModuleOp> candidate;
      LogicalResult realized = failure();
      {
        ScopedCompilerOutputSuppression silence;
        realized = realizeCandidate(checkpoint, prepared->candidates[index], candidate, capturedReport, summary);
      }
      PimCandidateResult result = PimCandidateResult::Failure;
      if (failed(realized)) {
        summary["reason"] = "selected Spatial schedule failed deterministic realization";
        summary["status"] = "realization-failed";
      } else {
        if (pimDisableSynchronization)
          summary["synchronization_ablation_cost"] = analyzePimSpatialSynchronizationAblation(*candidate);
        ++lowered;
        result = lowerPimPipelineCandidate(*candidate, resources, summary);
      }
      realizationSeconds += summary.getNumber("realization_seconds").value_or(0);
      backendSeconds += summary.getNumber("backend_seconds").value_or(0);
      if (result == PimCandidateResult::Success) {
        summary["status"] = "lowered";
        winner = std::move(candidate);
        scheduleReport = std::move(capturedReport);
      } else if (result == PimCandidateResult::ResourceLimit)
        llvm::errs() << "Spatial pipeline candidate " << index
          << " rejected by a backend resource limit; rejection recorded in pipeline_selection.json\n";
      attempts.push_back(std::move(summary));
      return result;
    };
    std::string reason;
    std::optional<size_t> selected;
    if (single) {
      if (attempt(0) == PimCandidateResult::Success) selected = 0;
      reason = selected ? "one distinct schedule; no cost comparison required"
        : "the single schedule failed backend lowering; no alternatives are available";
    } else selected = lowerSelectedPimCandidate(costs, attempt, reason);
    report["backend_attempts"] = std::move(attempts);
    report["lowered_candidates"] = static_cast<int64_t>(lowered);
    report["costed_candidates"] = single ? 0 : static_cast<int64_t>(costs.size());
    size_t attemptCount = report.getArray("backend_attempts")->size();
    report["backend_retries"] = static_cast<int64_t>(attemptCount ? attemptCount - 1 : 0);
    report["winner_realization_seconds"] = realizationSeconds;
    report["backend_lowering_seconds"] = backendSeconds;
    if (!selected) {
      report["status"] = "backend-failure";
      report["reason"] = reason;
      std::string failures;
      for (auto &entry : *report.getArray("backend_attempts")) {
        auto *failed = entry.getAsObject();
        failures += "\ncandidate " + std::to_string(*failed->getInteger("candidate"))
          + ": " + failed->getString("reason").value_or("backend failure").str();
      }
      checkpoint.emitError("Spatial pipeline backend selection failed; ") << reason << failures;
      writeSelectionReport(std::move(report)); signalPassFailure(); return;
    }
    auto finalize = Clock::now();
    bool spatialArtifacts = exportStage != SpatialDataflowExportStage::None
      || pimPeriodicDiagnostics || isPimIrDumpEnabled();
    if (spatialArtifacts) {
      if (failed(realizeSpatialSchedule(checkpoint, prepared->candidates[*selected], target,
            stages, exportStage, pimPeriodicDiagnostics ? &prepared->graph : nullptr))) {
        signalPassFailure(); return;
      }
    } else {
      auto file = openDialectDumpFileWithExtension("spatial3_scheduled_no_comm", "/reports", "txt");
      if (file) file << scheduleReport;
    }
    checkpoint->setAttrs((*winner)->getAttrs());
    checkpoint.getBodyRegion().takeBody(winner->getBodyRegion());
    report["selection"] = static_cast<int64_t>(*selected);
    report["status"] = "selected";
    report["reason"] = reason;
    report["finalization_seconds"] = elapsed(finalize);
    report["total_seconds"] = elapsed(start);
    llvm::errs() << "Spatial queue-aware pipeline: candidate " << *selected << " of " << prepared->candidates.size()
                 << "; " << reason << "\n";
    writeSelectionReport(std::move(report));
  }
};
} // namespace

PimCandidateResult onnx_mlir::lowerPimPipelineCandidate(ModuleOp module,
    const SpatialTargetResources &resources, llvm::json::Object &summary) {
  bool classified = true, anyError = false;
  std::string diagnostic;
  size_t count = 0;
  ScopedDiagnosticHandler handler(module.getContext(), [&](Diagnostic &message) {
    if (message.getSeverity() == DiagnosticSeverity::Error) {
      anyError = true;
      classified &= pim::isCandidateResourceError(message);
      if (++count <= 8) diagnostic += message.str().substr(0, 2048) + "\n";
    }
    return success();
  });
  auto start = Clock::now();
  PassManager pm(module.getContext());
  addPimLoweringPasses(pm, resources, EmitPimCodegen, /*reportProgress=*/false);
  if (enableTiming) pm.addInstrumentation(std::make_unique<CandidateTiming>(summary));
  LogicalResult result = pm.run(module);
  summary["backend_seconds"] = elapsed(start);
  if (succeeded(result)) return PimCandidateResult::Success;
  summary["reason"] = diagnostic.empty() ? "backend failed without a classified diagnostic" : diagnostic;
  bool resource = anyError && classified;
  summary["status"] = resource ? "resource-limit" : "systemic-or-unclassified-error";
  return resource ? PimCandidateResult::ResourceLimit : PimCandidateResult::Failure;
}

std::optional<size_t> onnx_mlir::lowerSelectedPimCandidate(llvm::ArrayRef<WholeScheduleCost> costs,
    const std::function<PimCandidateResult(size_t)> &lower, std::string &reason) {
  if (costs.empty()) {
    reason = "Spatial pipeline requires at least one candidate schedule";
    return std::nullopt;
  }
  for (const auto &cost : costs)
    if (!cost.eligible()) {
      reason = "cannot compare the complete candidate set: unavailable cost or unproved eligibility";
      return std::nullopt;
    }
  std::vector<WholeScheduleCost> remaining(costs.begin(), costs.end());
  std::vector<size_t> indices;
  for (size_t i = 0; i < costs.size(); ++i) indices.push_back(i);
  while (!indices.empty()) {
    size_t selected = selectPipelineCandidate(remaining, reason);
    auto result = lower(indices[selected]);
    if (result == PimCandidateResult::Success) return indices[selected];
    if (result != PimCandidateResult::ResourceLimit) {
      reason = "systemic or unclassified backend failure; recovery is not justified";
      return std::nullopt;
    }
    remaining.erase(remaining.begin() + selected);
    indices.erase(indices.begin() + selected);
  }
  reason = "all candidate schedules failed backend resource constraints";
  return std::nullopt;
}

std::unique_ptr<Pass> onnx_mlir::createPimPipelineSelectionPass(
    const SchedulingTarget &target, const SpatialTargetResources &resources,
    size_t stages, SpatialDataflowExportStage exportStage) {
  return std::make_unique<PimPipelineSelectionPass>(target, resources, stages, exportStage);
}