#!/usr/bin/env ruby

require "digest"
require "json"
require "optparse"
require "pathname"
require "yaml"

TYPE_BYTES = {
  "int32" => 4,
  "bfloat16" => 2,
  "float32" => 4,
  "int8" => 1,
  "uint8" => 1,
  "PrimitiveType.S32" => 4,
  "PrimitiveType.BF16" => 2,
  "PrimitiveType.F32" => 4,
  "PrimitiveType.S8" => 1,
  "PrimitiveType.U8" => 1,
  "PrimitiveType.U16" => 2,
  "PrimitiveType.S16" => 2,
  "PrimitiveType.U32" => 4,
}.freeze

options = {
  output: nil,
  full_hash: false,
  hash_limit: 64 * 1024 * 1024,
}

OptionParser.new do |parser|
  parser.banner = "usage: calculet_package_validate.rb [options] MODEL_DIR"
  parser.on("-o", "--output FILE", "write JSON report to FILE") { |v| options[:output] = v }
  parser.on("--full-hash", "SHA-256 every file, including large parameter blocks") { options[:full_hash] = true }
  parser.on("--hash-limit-mib N", Integer, "fast-mode hash limit (default: 64 MiB)") do |v|
    options[:hash_limit] = v * 1024 * 1024
  end
end.parse!

abort "MODEL_DIR is required" unless ARGV.length == 1

root = Pathname.new(ARGV.first).expand_path
abort "not a directory: #{root}" unless root.directory?

findings = []

add_finding = lambda do |severity, code, message, evidence = nil|
  item = { severity: severity, code: code, message: message }
  item[:evidence] = evidence if evidence
  findings << item
end

safe_yaml = lambda do |path|
  YAML.safe_load(path.read, permitted_classes: [], permitted_symbols: [], aliases: true)
rescue Psych::Exception => e
  add_finding.call("error", "YAML_PARSE", e.message, path.relative_path_from(root).to_s)
  nil
end

relative = lambda do |path|
  Pathname.new(path).expand_path.relative_path_from(root).to_s
rescue ArgumentError
  path.to_s
end

files = root.glob("**/*", File::FNM_DOTMATCH).select(&:file?)
symlinks = root.glob("**/*", File::FNM_DOTMATCH).select(&:symlink?)

symlinks.each do |link|
  begin
    resolved = link.realpath
    unless resolved.to_s.start_with?(root.realpath.to_s + File::SEPARATOR)
      add_finding.call("error", "SYMLINK_ESCAPE", "symlink resolves outside package", relative.call(link))
    end
  rescue Errno::ENOENT
    add_finding.call("error", "BROKEN_SYMLINK", "broken symlink", relative.call(link))
  end
end

required_exact = %w[hparam.yaml model_memory_reserved_info.txt param_blk0.bin param_blk1.bin]
required_exact.each do |name|
  path = root.join(name)
  if !path.file?
    add_finding.call("error", "REQUIRED_FILE", "required file is missing", name)
  elsif path.size.zero?
    add_finding.call("error", "EMPTY_FILE", "required file is empty", name)
  end
end

manifest_candidates = root.glob("*.yaml").reject do |path|
  path.basename.to_s == "hparam.yaml" || path.basename.to_s.end_with?("_op_io_buf_ptr.yaml")
end

if manifest_candidates.length != 1
  add_finding.call("error", "ROOT_MANIFEST_COUNT",
                   "expected exactly one root model manifest, got #{manifest_candidates.length}")
end

manifest_path = manifest_candidates.first
manifest = manifest_path ? safe_yaml.call(manifest_path) : nil

model_summary = {}
submodels = []
supply_chain = {}

if manifest.is_a?(Hash)
  cfg = manifest["one_model_cfg"]
  unless cfg.is_a?(Hash)
    add_finding.call("error", "MODEL_CFG", "one_model_cfg must be a mapping", relative.call(manifest_path))
    cfg = {}
  end

  cfg.each do |logical_name, item|
    unless item.is_a?(Hash)
      add_finding.call("error", "SUBMODEL_CFG", "submodel config must be a mapping", logical_name)
      next
    end

    info = item["model_info"] || {}
    device = item["device_cfg"] || {}
    attrs = item["model_attrs"].to_s
    kind = attrs.empty? ? (logical_name.include?("prefill") ? "prefill" : "decode") : attrs
    submodels << {
      logical_name: logical_name,
      kind: kind,
      model_path: item["model_path"],
      batch: info["batch_num"],
      max_seq: info["max_seq_len"],
      chips: device["chip_nums"],
      calcores_per_chip: device["calcore_num"],
      input_shape: item["input_shape"],
      flash_attention: item.dig("calcc_env_cfg", "enable_flash_attention"),
      one_token: item.dig("calcc_env_cfg", "enable_one_token"),
      dynamic_d2d: item.dig("calcc_env_cfg", "dynamic_d2d"),
    }
  end

  kinds = submodels.map { |s| s[:kind] }
  %w[prefill decode].each do |kind|
    count = kinds.count(kind)
    add_finding.call("error", "SUBMODEL_KIND_COUNT", "expected one #{kind}, got #{count}") unless count == 1
  end

  batches = submodels.map { |s| s[:batch] }.compact.uniq
  seqs = submodels.map { |s| s[:max_seq] }.compact.uniq
  chips = submodels.map { |s| s[:chips] }.compact.uniq
  add_finding.call("error", "BATCH_MISMATCH", "submodel batch values differ: #{batches}") if batches.length > 1
  add_finding.call("error", "SEQ_MISMATCH", "submodel max_seq values differ: #{seqs}") if seqs.length > 1
  add_finding.call("error", "CHIP_MISMATCH", "submodel chip counts differ: #{chips}") if chips.length > 1

  model_summary = {
    name: manifest["model_name"],
    seq_len: manifest["seq_len"],
    batch: batches.first,
    max_seq: seqs.first,
    chips: chips.first,
  }
  if model_summary[:seq_len] && model_summary[:max_seq] && model_summary[:seq_len] != model_summary[:max_seq]
    add_finding.call("error", "ROOT_SEQ_MISMATCH", "root seq_len differs from submodel max_seq")
  end

  supply_chain = {
    onnx_hash: manifest["onnx_hash"],
    compiler_commits: manifest["git_commit_id"],
    golden_hash: manifest["golden_hash"],
  }
  if Array(supply_chain[:onnx_hash]).empty?
    add_finding.call("blocked_reproducibility", "ONNX_HASH_MISSING", "ONNX hash is missing")
  end
end

memory = { parameter_files: {}, reservations: {}, parameter_maps: [] }
memory_path = root.join("model_memory_reserved_info.txt")
if memory_path.file?
  memory_path.each_line.with_index(1) do |line, line_no|
    case line
    when /\A(dram|sram|sync)_rsvd\s+vaddr\s+(0x[0-9a-fA-F]+)\s+size\s+(\d+)/
      memory[:reservations][Regexp.last_match(1)] = {
        vaddr: Regexp.last_match(2), bytes: Regexp.last_match(3).to_i,
      }
    when /\A([^\s]+\.bin)\s+offset\s+(0x[0-9a-fA-F]+)\s+vaddr\s+(0x[0-9a-fA-F]+)\s+size\s+(0x[0-9a-fA-F]+)\s+chip\s+([0-9a-fA-F]+)\s+([0-9a-fA-F]+)\s+([0-9a-fA-F]+)/
      file_name = Regexp.last_match(1)
      offset = Regexp.last_match(2).to_i(16)
      vaddr = Regexp.last_match(3).to_i(16)
      size = Regexp.last_match(4).to_i(16)
      mask = [Regexp.last_match(5), Regexp.last_match(6), Regexp.last_match(7)].join
      target = root.join(file_name)
      if !target.file?
        add_finding.call("error", "PARAM_FILE_MISSING", "mapped parameter file missing", file_name)
      elsif offset > target.size || size > target.size - offset
        add_finding.call("error", "PARAM_RANGE", "parameter mapping exceeds file", "#{file_name}:#{line_no}")
      end
      add_finding.call("error", "CHIP_MASK_ZERO", "parameter map has zero chip mask", "line #{line_no}") if mask.to_i(16).zero?
      memory[:parameter_maps] << {
        file: file_name, offset: offset, vaddr: vaddr, bytes: size, chip_mask: mask,
      }
    when /\Adevice_memory_required:\s+(\d+)\s+MB/
      memory[:device_memory_required_raw] = { value: Regexp.last_match(1).to_i, unit: "MB" }
      add_finding.call("warning_vendor", "DEVICE_MEMORY_UNIT",
                       "device_memory_required unit/value is not credible for capacity decisions", "line #{line_no}")
    end
  end
end

%w[param_blk0.bin param_blk1.bin].each do |name|
  path = root.join(name)
  memory[:parameter_files][name] = path.size if path.file?
end

metadata = []
root.glob("*/submodel_memory_reserved_info.txt").sort.each do |path|
  item = { path: relative.call(path), inputs: [], outputs: [], csrs: [], files: [] }
  path.each_line.with_index(1) do |line, line_no|
    case line
    when /\Asmodel_type\s+(\S+)/
      item[:smodel_type] = Regexp.last_match(1)
    when /\An_batch\s+(\d+)/
      item[:batch] = Regexp.last_match(1).to_i
    when /\Aicsr\s+(.+?)\s+offset\s+(0x[0-9a-fA-F]+)/
      item[:csrs] << { name: Regexp.last_match(1), offset: Regexp.last_match(2) }
    when /\A(ibuf|obuf)\s+(.+?)\s+shape\s+\[([^\]]+)\]\s+type\s+(\S+)\s+vaddr\s+(0x[0-9a-fA-F]+)\s+(0x[0-9a-fA-F]+)\s+size\s+(\d+)/
      io_kind = Regexp.last_match(1)
      dims = Regexp.last_match(3).split(",").map { |v| Integer(v.strip) }
      dtype = Regexp.last_match(4)
      bytes = Regexp.last_match(7).to_i
      tensor = {
        name: Regexp.last_match(2), shape: dims, dtype: dtype,
        ping_vaddr: Regexp.last_match(5), pong_vaddr: Regexp.last_match(6), bytes: bytes,
      }
      if TYPE_BYTES.key?(dtype)
        expected = dims.inject(1, :*) * TYPE_BYTES.fetch(dtype)
        if expected != bytes
          add_finding.call("error", "TENSOR_SIZE", "tensor size formula mismatch",
                           "#{relative.call(path)}:#{line_no} expected=#{expected} actual=#{bytes}")
        end
      else
        add_finding.call("error", "TENSOR_DTYPE", "unknown tensor dtype #{dtype}", "#{relative.call(path)}:#{line_no}")
      end
      (io_kind == "ibuf" ? item[:inputs] : item[:outputs]) << tensor
    when /\A(\S+\.(?:bin|elf|so))\b.*\bchip\s+/
      file_name = Regexp.last_match(1)
      item[:files] << file_name
      declared = path.dirname.join(file_name)
      add_finding.call("error", "DECLARED_FILE", "declared submodel file missing", relative.call(declared)) unless declared.file?
    end
  end
  if path.dirname.basename.to_s.include?("prefill") && item[:smodel_type] == "llm_decode"
    add_finding.call("warning_vendor", "PREFILL_MODEL_TYPE", "prefill metadata is marked llm_decode", relative.call(path))
  end
  metadata << item
end

if metadata.length != 2
  add_finding.call("error", "SUBMODEL_METADATA_COUNT", "expected two submodel metadata files, got #{metadata.length}")
end

graphs = {}
root.glob("*_op_io_buf_ptr.yaml").sort.each do |path|
  graph = safe_yaml.call(path)
  next unless graph.is_a?(Hash)

  kind = path.basename.to_s.include?("prefill") ? "prefill" : "decode"
  histogram = Hash.new(0)
  unknown_dtypes = []
  graph.each do |op_name, op|
    histogram[op_name.sub(/_\d+\z/, "")] += 1
    next unless op.is_a?(Hash)
    op.each_value do |tensor|
      next unless tensor.is_a?(Array) && tensor.length >= 3
      dtype = tensor[2].to_s
      unknown_dtypes << dtype unless TYPE_BYTES.key?(dtype)
    end
  end
  unknown_dtypes.uniq.each do |dtype|
    add_finding.call("error", "GRAPH_DTYPE", "unknown graph dtype #{dtype}", relative.call(path))
  end
  graphs[kind] = { path: relative.call(path), op_count: graph.length, op_histogram: histogram.sort.to_h }

  prefix = path.basename.to_s.sub("_op_io_buf_ptr.yaml", "")
  root.glob("#{prefix}.chip[01].profparts").each do |prof|
    unknown = prof.each_line.map(&:strip).reject { |name| name.empty? || name == "tail" }.reject { |name| graph.key?(name) }
    unless unknown.empty?
      add_finding.call("error", "PROFPART_UNKNOWN_OP", "profparts references unknown operators",
                       "#{relative.call(prof)} first=#{unknown.first} count=#{unknown.length}")
    end
  end
end

%w[prefill decode].each do |kind|
  add_finding.call("error", "GRAPH_MISSING", "#{kind} operator graph is missing") unless graphs.key?(kind)
end

hashes = {}
files.sort.each do |path|
  rel = relative.call(path)
  if options[:full_hash] || path.size <= options[:hash_limit]
    hashes[rel] = { bytes: path.size, sha256: Digest::SHA256.file(path).hexdigest }
  else
    hashes[rel] = { bytes: path.size, sha256: nil, skipped: "larger_than_fast_mode_limit" }
  end
end

severity_order = {
  "error" => 0,
  "blocked_reproducibility" => 1,
  "warning_vendor" => 2,
  "warning_service" => 3,
  "info" => 4,
}
findings.sort_by! { |f| [severity_order.fetch(f[:severity], 99), f[:code], f[:evidence].to_s] }

verdict = if findings.any? { |f| f[:severity] == "error" }
            "fail"
          elsif findings.any? { |f| f[:severity] != "info" }
            "pass_with_warnings"
          else
            "pass"
          end

report = {
  schema: "calculet-package-validation/v1",
  package_root: root.to_s,
  generated_at_utc: Time.now.utc.strftime("%Y-%m-%dT%H:%M:%SZ"),
  model: model_summary,
  submodels: submodels,
  runtime_metadata: metadata,
  memory: memory,
  graph: graphs,
  supply_chain: supply_chain,
  inventory: { file_count: files.length, symlink_count: symlinks.length, files: hashes },
  findings: findings,
  verdict: verdict,
}

json = JSON.pretty_generate(report) + "\n"
if options[:output]
  File.write(options[:output], json)
else
  puts json
end

exit(verdict == "fail" ? 1 : 0)