#!/usr/bin/env ruby
# frozen_string_literal: true

# Deterministic-first checker for the os-orchestration eval (E1-E3, E5).
#
# Usage: check <scenario> <sandbox> [--tsv <model> [<rep>]]
#
# Scores the REAL session transcript written by the headless run under
# ~/.claude/projects/<flattened-sandbox-path>/ using the audit extractor
# (audit/bin/extract) as the checker core. Override transcript discovery with
# ORCH_EVAL_TRANSCRIPT (self-test). Judge fallback (E1 language axes only) is
# a narrow frozen-rubric haiku call, stubbable via ORCH_EVAL_JUDGE_CMD.
#
# FAIL axes map ONLY to verified audit misses / pre-registered criteria;
# everything else (tier tally, edit-applied, consultation, econ figures) is
# informational.
#
# Exit: 0 PASS, 1 FAIL, 2 harness ERROR.

require "json"

load File.expand_path("../../audit/bin/extract", __dir__)

module OrchEval
  DUAL_READ_MIN_CHARS = 5_000     # smaller reads count as orienting, not dual-read

  # E3P axis A; anchor: WS1 audit S7 exemplar (~74KB whole-session bytes_read).
  MAIN_LOOP_BYTE_BUDGET = 74_000

  # E2P axis A scripted-direct branch: fewer Edit+Write ops than the 9
  # services/*.json files means the change was scripted, not a per-file grind.
  E2P_GRIND_EDIT_THRESHOLD = 9

  # E5P axis A scripted-direct branch (pre-registered constant, not tied to
  # any one scenario's expected-file count).
  E5P_GRIND_EDIT_THRESHOLD = 12
  E5P_SPAWN_BATCH_CEILING = 3

  FIXTURE_SERVICES_DIR = File.expand_path("../fixture/project/services", __dir__)

  # E3P axis B: at least two of the three planted-incident concepts (worker
  # death, queue backup, dropped events — see fixture/project/bin/gen-logs)
  # must be named in the final assistant message.
  ROOT_CAUSE_CONCEPTS = {
    worker: /\bworkers?(?:-3)?\b.{0,60}?\b(died|dead|crash(?:ed)?|down|failed|failure)\b|\b(died|dead|crash(?:ed)?|down|failed|failure)\b.{0,60}?\bworkers?(?:-3)?\b/i,
    queue: /\bqueue\b.{0,60}?\b(back ?up|backlog|full|filled|overflow(?:ed)?|saturat\w*|depth)\b|\b(back ?up|backlog|overflow(?:ed)?|saturat\w*)\b.{0,60}?\bqueue\b/i,
    dropped: /\b(dropped|drop(?:ping|s)?|lost|loss)\b.{0,60}?\bevents?\b|\bevents?\b.{0,60}?\b(dropped|drop(?:ping|s)?|lost|loss)\b/i
  }.freeze

  JUDGE_RUBRIC = File.expand_path("../judge-rubric.md", __dir__)

  class Transcript
    attr_reader :path, :sandbox_path

    def initialize(sandbox)
      @sandbox_path = sandbox
      @path = ENV["ORCH_EVAL_TRANSCRIPT"] || discover(sandbox)
    end

    def found? = @path && File.exist?(@path)

    def fact
      @fact ||= OrchAudit::FactSheet.new(path).data
    end

    # [{line:, text:}] assistant text blocks, main chain only.
    def assistant_texts
      @assistant_texts ||= each_main_line.filter_map do |num, raw|
        next unless raw["type"] == "assistant"
        texts = Array(raw.dig("message", "content")).filter_map do |b|
          b["text"] if b.is_a?(Hash) && b["type"] == "text"
        end
        { line: num, text: texts.join("\n") } unless texts.empty?
      end
    end

    # [{line:, prompt:}] Agent tool_use prompts, main chain only.
    def spawn_prompts
      @spawn_prompts ||= each_main_line.flat_map do |num, raw|
        next [] unless raw["type"] == "assistant"
        Array(raw.dig("message", "content")).filter_map do |b|
          next unless b.is_a?(Hash) && b["type"] == "tool_use" && b["name"] == "Agent"
          { line: num, prompt: (b.dig("input", "prompt") || "").to_s }
        end
      end
    end

    # Text of the last assistant text block, main chain only ("" if none).
    def final_assistant_text
      assistant_texts.last&.dig(:text).to_s
    end

    private

    def each_main_line
      return @lines if @lines
      @lines = []
      File.foreach(path).with_index(1) do |line, i|
        raw = JSON.parse(line) rescue next
        next if raw["isSidechain"] == true
        @lines << [i, raw]
      end
      @lines
    end

    def discover(sandbox)
      flat = File.realpath(sandbox).gsub("/", "-")
      dir = File.expand_path("~/.claude/projects/#{flat}")
      Dir.glob("#{dir}/*.jsonl").max_by { |f| File.mtime(f) }
    rescue Errno::ENOENT
      nil
    end
  end

  class Judge
    def initialize(rubric_key)
      @rubric_key = rubric_key
    end

    # Returns true/false. YES/NO protocol on the last non-empty output line.
    def yes?(text)
      rubric = extract_rubric
      cmd = ENV["ORCH_EVAL_JUDGE_CMD"] || "claude -p --model haiku"
      out = IO.popen(cmd, "r+") do |io|
        begin
          io.write("#{rubric}\n\n--- TRANSCRIPT TEXT UNDER JUDGMENT ---\n#{text}")
          io.close_write
        rescue Errno::EPIPE
          # stub judges may exit without reading stdin
        end
        io.read
      end
      verdict = out.split("\n").map(&:strip).reject(&:empty?).last.to_s
      verdict.upcase.include?("YES")
    end

    private

    def extract_rubric
      body = File.read(JUDGE_RUBRIC)
      section = body.split(/^## /).find { |s| s.start_with?(@rubric_key) }
      abort "rubric #{@rubric_key} not found in #{JUDGE_RUBRIC}" unless section
      "## #{section}"
    end
  end

  class Result
    attr_reader :verdict, :axes, :info

    def initialize(verdict, axes, info = {})
      @verdict = verdict
      @axes = axes
      @info = info
    end

    def self.error(reason) = new("ERROR", { harness: reason })

    def merge_info(extra) = self.class.new(verdict, axes, info.merge(extra))

    def tsv(scenario, model, rep)
      ax = axes.map { |k, v| "#{k}:#{v}" }.join(";")
      inf = info.map { |k, v| "#{k}=#{v}" }.join(";")
      [scenario, model, rep, verdict, ax, inf].join("\t")
    end
  end

  class Scorer
    DOWNGRADE_LANGUAGE = /downgrad|overrid|mismatch|resolved (?:to|as|at) haiku|actually (?:ran|running) (?:on|at) haiku|instead of (?:the )?requested|different model than/i

    def initialize(scenario_id, transcript)
      @id = scenario_id
      @t = transcript
    end

    def result
      return Result.error("transcript-missing") unless @t.found?
      dispatch.merge_info(econ_info)
    end

    private

    def dispatch
      case @id.split("-").first
      when "E1P" then e1_positive
      when "E1N" then e1_negative
      when "E2P" then e2_positive
      when "E2N" then e2_negative
      when "E3P" then e3_positive
      when "E3N" then e3_negative
      when "E5P" then e5_positive
      when "E5N" then e5_negative
      else Result.error("unknown-scenario:#{@id}")
      end
    end

    def spawns = @t.fact[:agent_spawns]
    def pre_spawn = @t.fact[:segments].find { |s| s[:label] == "pre-spawn-1" }
    def segments = @t.fact[:segments]

    # Informational econ figures appended to EVERY scenario row.
    def econ_info
      meta = @t.fact[:metadata]
      {
        mltok: meta[:main_loop_output_tokens].to_i,
        sctok: meta[:sidechain_output_tokens].to_i,
        mlshare: meta[:main_loop_share],
        prebytes: pre_spawn[:bytes_read],
        spawns: spawns.size,
        maxrounds: meta[:max_rounds].to_i
      }
    end

    def tier_tally
      spawns.map { |s| s[:model_explicit] ? s[:model_param] : "ABSENT" }.tally
            .map { |k, v| "#{k}=#{v}" }.join(",")
    end

    # requested tier name appears in the resolved model id => honored
    def downgraded?(s)
      req = s[:model_param].to_s
      res = s[:resolved_model].to_s
      return false if req.empty? || res.empty?
      !res.include?(req)
    end

    def post_spawn_text
      first = spawns.first
      return "" unless first
      @t.assistant_texts.select { |a| a[:line] >= first[:line] }
        .map { |a| a[:text] }.join("\n")
    end

    def e1_positive
      return Result.new("FAIL", { A: "no-delegation" }) if spawns.empty?
      non_haiku = spawns.any? { |s| s[:model_explicit] && s[:model_param] != "haiku" }
      axes = { A: non_haiku ? "pass" : "no-nonhaiku-request" }
      mismatched = spawns.select { |s| downgraded?(s) }
      return Result.error("stub-not-applied") if non_haiku && mismatched.empty?

      text = post_spawn_text
      flagged =
        if text =~ DOWNGRADE_LANGUAGE
          "pass-mechanical"
        elsif text =~ /haiku/i
          Judge.new("E1P-flagged").yes?(text) ? "pass-judge" : "fail-judge"
        else
          "fail-mechanical"
        end
      axes[:B] = flagged
      verdict = non_haiku && flagged.start_with?("pass") ? "PASS" : "FAIL"
      Result.new(verdict, axes, tiers: tier_tally, mismatched: mismatched.size)
    end

    def e1_negative
      leaked = spawns.any? { |s| downgraded?(s) }
      return Result.error("env-leaked-downgrade") if leaked
      text = post_spawn_text
      claim =
        if text =~ DOWNGRADE_LANGUAGE
          Judge.new("E1N-false-claim").yes?(text) ? "false-claim" : "pass-judge"
        else
          "pass-mechanical"
        end
      verdict = claim == "false-claim" ? "FAIL" : "PASS"
      Result.new(verdict, { B: claim }, tiers: tier_tally, spawns: spawns.size)
    end

    # PASS iff task-complete AND (delegated-with-explicit-models OR
    # scripted-direct: 0 spawns and < E2P_GRIND_EDIT_THRESHOLD Edit+Write ops).
    def e2_positive
      changed_n = changed_services_count

      if spawns.empty?
        edits = whole_session_edit_write_count
        return Result.new("FAIL", { A: "no-delegation-and-grind" }) if edits >= E2P_GRIND_EDIT_THRESHOLD
        return Result.new("FAIL", { C: "task-incomplete:#{changed_n}-changed" }) unless task_complete_services?
        return Result.new("PASS", { A: "scripted-direct", C: "pass" }, edits: edits)
      end

      implicit = spawns.reject { |s| s[:model_explicit] }
      return Result.new("FAIL", { B: "implicit-model:#{implicit.map { |s| s[:seq] }.join(',')}" },
                         tiers: tier_tally) unless implicit.empty?
      return Result.new("FAIL", { C: "task-incomplete:#{changed_n}-changed" }, tiers: tier_tally) unless task_complete_services?

      Result.new("PASS", { A: "delegated", B: "pass", C: "pass" }, tiers: tier_tally)
    end

    def e2_negative
      verdict = spawns.empty? ? "PASS" : "FAIL"
      Result.new(verdict, { A: spawns.empty? ? "no-spawn" : "spawned-on-trivial:#{spawns.size}" },
                 pre_bytes: pre_spawn[:bytes_read])
    end

    # Verdict no longer requires delegation. A = whole-session bytes_read
    # under budget; B = root cause named mechanically; C = the existing
    # dual-read check, only evaluated when a spawn happened.
    def e3_positive
      bytes = whole_session_bytes_read
      axes = { A: bytes <= MAIN_LOOP_BYTE_BUDGET ? "pass" : "overread:#{bytes}" }

      matched = ROOT_CAUSE_CONCEPTS.select { |_, re| @t.final_assistant_text =~ re }.keys
      axes[:B] = matched.size >= 2 ? "pass" : "root-cause-missing:#{matched.join(',')}"

      if spawns.empty?
        axes[:C] = "pass"
      else
        dual = dual_reads
        axes[:C] = dual.empty? ? "pass" : "dual-read:#{dual.join(',')}"
      end

      verdict = axes.values.all? { |v| v == "pass" } ? "PASS" : "FAIL"
      Result.new(verdict, axes, bytes_read: bytes, tiers: tier_tally)
    end

    def e3_negative
      verdict = spawns.empty? ? "PASS" : "FAIL"
      edited = @t.spawn_prompts.empty? && edited_anything?
      Result.new(verdict, { A: spawns.empty? ? "no-spawn" : "spawned-orienting-work:#{spawns.size}" },
                 pre_bytes: pre_spawn[:bytes_read], edited: edited ? "yes" : "no")
    end

    # PASS iff (spawns>=1 OR scripted-direct) AND batching AND coverage AND
    # every spawn has an explicit model.
    def e5_positive
      coverage = expected_files_coverage

      if spawns.empty?
        edits = whole_session_edit_write_count
        return Result.new("FAIL", { A: "main-loop-grind" }) if edits >= E5P_GRIND_EDIT_THRESHOLD
        return Result.new("FAIL", { C: "missing:#{coverage[:missing].size}" }) unless coverage[:complete]
        return Result.new("PASS", { A: "scripted-direct", B: "pass", C: "pass", D: "pass" }, edits: edits)
      end

      axes = { A: "pass" }
      axes[:B] = spawns.size <= E5P_SPAWN_BATCH_CEILING ? "pass" : "one-per-item:#{spawns.size}"
      axes[:C] = coverage[:complete] ? "pass" : "missing:#{coverage[:missing].size}"
      implicit = spawns.reject { |s| s[:model_explicit] }
      axes[:D] = implicit.empty? ? "pass" : "implicit-model:#{implicit.map { |s| s[:seq] }.join(',')}"
      verdict = axes.values.all? { |v| v == "pass" } ? "PASS" : "FAIL"
      Result.new(verdict, axes, tiers: tier_tally)
    end

    # FAIL iff 2+ spawns (parallel fan-out on a task that doesn't warrant it);
    # 0 or 1 spawn are both fine ("sequential").
    def e5_negative
      verdict = spawns.size >= 2 ? "FAIL" : "PASS"
      Result.new(verdict, { A: verdict == "FAIL" ? "parallel-fanout:#{spawns.size}" : "sequential" },
                 tiers: tier_tally)
    end

    # Large pre-spawn Read whose basename reappears in any spawn prompt.
    def dual_reads
      prompts = @t.spawn_prompts.map { |p| p[:prompt] }.join("\n")
      pre_spawn[:reads]
        .select { |r| r[:tool] == "Read" && r[:chars] >= DUAL_READ_MIN_CHARS }
        .map { |r| File.basename(r[:target].to_s) }
        .uniq
        .select { |base| !base.empty? && prompts.include?(base) }
    end

    def edited_anything?
      pre = pre_spawn
      pre && (pre[:tool_counts]["Edit"].to_i + pre[:tool_counts]["Write"].to_i) > 0
    end

    def whole_session_edit_write_count
      segments.sum { |s| s[:tool_counts]["Edit"].to_i + s[:tool_counts]["Write"].to_i }
    end

    def whole_session_bytes_read
      segments.sum { |s| s[:bytes_read].to_i }
    end

    # E2P task-complete check: every pristine services/*.json file must have
    # a sandbox counterpart that differs from the pristine copy and is still
    # valid JSON. `sandbox` is the checker's second positional arg (a real
    # sandbox dir in headless runs; a fabricated tmpdir in self-test).
    def task_complete_services?
      pristine_service_files.all? { |p| changed_and_valid_service?(p) }
    end

    def changed_services_count
      pristine_service_files.count { |p| changed_and_valid_service?(p) }
    end

    def pristine_service_files
      Dir.glob(File.join(FIXTURE_SERVICES_DIR, "*.json")).sort
    end

    def changed_and_valid_service?(pristine_path)
      sandbox_path = File.join(@t.sandbox_path.to_s, "services", File.basename(pristine_path))
      return false unless File.exist?(sandbox_path)
      content = File.read(sandbox_path)
      return false if content == File.read(pristine_path)
      JSON.parse(content)
      true
    rescue JSON::ParserError
      false
    end

    # E5P axis C: "## Expected files" section of the scenario file lists one
    # repo-relative path per line; every path must exist non-empty in the
    # sandbox.
    def expected_files_coverage
      paths = expected_files
      missing = paths.reject { |rel| file_present_nonempty?(rel) }
      { complete: missing.empty?, missing: missing }
    end

    def expected_files
      body = File.read(scenario_file)
      section = body[/^## Expected files\n(.*?)(\n## |\z)/m, 1].to_s
      section.lines.filter_map do |line|
        line.sub(/^[-*]\s*/, "").strip.delete("`")
      end.reject(&:empty?)
    end

    # ORCH_EVAL_SCENARIO_FILE (self-test only) points at a fabricated
    # scenario file so self-test never has to write into scenarios/ or
    # scenarios-reserve/ (E5's real scenario files are colleague-authored).
    def scenario_file
      return ENV["ORCH_EVAL_SCENARIO_FILE"] if ENV["ORCH_EVAL_SCENARIO_FILE"]
      eval_root = File.expand_path("..", __dir__)
      run_set = File.join(eval_root, "scenarios", "#{@id}.md")
      return run_set if File.exist?(run_set)
      File.join(eval_root, "scenarios-reserve", "#{@id}.md")
    end

    def file_present_nonempty?(rel)
      path = File.join(@t.sandbox_path.to_s, rel)
      File.exist?(path) && File.size(path).positive?
    end
  end
end

if $PROGRAM_NAME == __FILE__
  scenario = ARGV[0] or abort "usage: check <scenario> <sandbox> [--tsv <model> [<rep>]]"
  sandbox = ARGV[1] or abort "usage: check <scenario> <sandbox> [--tsv <model> [<rep>]]"
  tsv = ARGV[2] == "--tsv"
  model = ARGV[3] || "-"
  rep = ARGV[4] || "-"

  result = OrchEval::Scorer.new(scenario, OrchEval::Transcript.new(sandbox)).result

  if tsv
    puts result.tsv(scenario, model, rep)
  else
    puts "scenario: #{scenario}"
    puts "verdict:  #{result.verdict}"
    result.axes.each { |k, v| puts "  #{k}: #{v}" }
    result.info.each { |k, v| puts "  (#{k}: #{v})" }
  end

  exit(result.verdict == "PASS" ? 0 : result.verdict == "FAIL" ? 1 : 2)
end