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

# Reduce a Claude Code session transcript (.jsonl) to a delegation fact-sheet.
#
# Usage: extract [--json] [--run-threshold N] SESSION.jsonl [...]
#
# Dual-use: human/auditor-readable markdown by default; --json for the
# structured form (checker core for a future orchestration eval).

require "json"
require "optparse"

module OrchAudit
  READ_TOOLS = %w[Read Grep Glob Bash].freeze
  RUN_TOOLS = %w[Read Grep Glob Edit Write].freeze

  class Line
    attr_reader :number, :raw

    def initialize(number, raw)
      @number = number
      @raw = raw
    end

    def type = raw["type"]
    def sidechain? = raw["isSidechain"] == true
    def message = raw["message"] || {}
    def timestamp = raw["timestamp"]
    def cwd = raw["cwd"]
    def model = message["model"]

    def content_blocks
      c = message["content"]
      c.is_a?(Array) ? c.select { |b| b.is_a?(Hash) } : []
    end

    def tool_uses
      return [] unless type == "assistant"
      content_blocks.select { |b| b["type"] == "tool_use" }
    end

    def tool_results
      return [] unless type == "user"
      content_blocks.select { |b| b["type"] == "tool_result" }
    end

    def human_prompt?
      return false unless type == "user" && !sidechain?
      c = message["content"]
      return !c.lstrip.start_with?("<task-notification>") if c.is_a?(String)
      c.is_a?(Array) && c.none? { |b| b.is_a?(Hash) && b["type"] == "tool_result" }
    end

    def task_notification
      return nil unless type == "user"
      c = message["content"]
      c = c.filter_map { |b| b["text"] if b.is_a?(Hash) }.join if c.is_a?(Array)
      return nil unless c.is_a?(String) && c.include?("<task-notification>")
      { task_id: c[%r{<task-id>([^<]+)</task-id>}, 1], chars: c.length }
    end
  end

  class ToolCall
    attr_reader :line, :block
    attr_accessor :result_chars

    def initialize(line, block)
      @line = line
      @block = block
      @result_chars = 0
    end

    def id = block["id"]
    def name = block["name"]
    def input = block["input"] || {}

    def target
      input["file_path"] || input["path"] || input["pattern"] || input["command"] || input["prompt"]
    end
  end

  class AgentSpawn
    attr_reader :call, :index
    attr_accessor :agent_id, :resolved_model, :notification_chars

    def initialize(call, index)
      @call = call
      @index = index
      @notification_chars = 0
    end

    def to_h
      {
        seq: index,
        line: call.line.number,
        timestamp: call.line.timestamp,
        subagent_type: call.input["subagent_type"] || "general-purpose",
        model_param: call.input["model"],
        model_explicit: call.input.key?("model"),
        resolved_model: resolved_model,
        run_in_background: call.input["run_in_background"],
        prompt_chars: (call.input["prompt"] || "").length,
        description: call.input["description"],
        result_chars: [call.result_chars, notification_chars].max
      }
    end
  end

  class Segment
    attr_reader :label

    def initialize(label)
      @label = label
      @calls = []
    end

    def add(call) = @calls << call

    def to_h
      counts = Hash.new(0)
      bytes = 0
      @calls.each do |c|
        counts[c.name] += 1
        bytes += c.result_chars if READ_TOOLS.include?(c.name)
      end
      { label: label, tool_counts: counts, bytes_read: bytes, total_calls: @calls.size }
    end
  end

  class MissedDelegationScan
    def initialize(calls, threshold)
      @calls = calls
      @threshold = threshold
    end

    def candidates
      runs = []
      current = []
      @calls.each do |call|
        if extend_run?(current, call)
          current << call
        else
          runs << current if qualifying?(current)
          current = RUN_TOOLS.include?(call.name) ? [call] : []
        end
      end
      runs << current if qualifying?(current)
      runs.map { |r| describe(r) }
    end

    private

    def extend_run?(run, call)
      return false if run.empty?
      run.first.name == call.name && RUN_TOOLS.include?(call.name)
    end

    def qualifying?(run)
      return false if run.size < @threshold
      run.map(&:target).uniq.size >= @threshold
    end

    def describe(run)
      {
        tool: run.first.name,
        length: run.size,
        lines: [run.first.line.number, run.last.line.number],
        distinct_targets: run.map(&:target).uniq.size,
        targets_sample: run.map(&:target).uniq.first(6).map { |t| t.to_s[0, 90] }
      }
    end
  end

  class FactSheet
    def initialize(path, run_threshold: 4)
      @path = path
      @run_threshold = run_threshold
    end

    def data
      @data ||= build
    end

    def to_json_s = JSON.pretty_generate(data)

    def to_markdown
      d = data
      md = +"# Delegation fact-sheet: #{File.basename(@path)}\n\n"
      md << "- Transcript: `#{@path}`\n"
      d[:metadata].each { |k, v| md << "- #{k}: #{v}\n" }
      md << "\n## Agent spawns (#{d[:agent_spawns].size})\n\n"
      if d[:agent_spawns].empty?
        md << "None.\n"
      else
        md << "| # | line | type | model param | resolved model | bg | prompt chars | result chars | description |\n"
        md << "|---|------|------|-------------|----------------|----|--------------|--------------|-------------|\n"
        d[:agent_spawns].each do |a|
          model = a[:model_explicit] ? a[:model_param].to_s : "(ABSENT)"
          md << "| #{a[:seq]} | #{a[:line]} | #{a[:subagent_type]} | #{model} | #{a[:resolved_model]} | #{a[:run_in_background] ? 'y' : 'n'} | #{a[:prompt_chars]} | #{a[:result_chars]} | #{a[:description]} |\n"
        end
      end
      md << "\n## Orchestrator tool profile (segments split at each Agent spawn)\n\n"
      md << "| segment | calls | bytes read | per-tool counts |\n|---------|-------|------------|------------------|\n"
      d[:segments].each do |s|
        counts = s[:tool_counts].map { |k, v| "#{k}:#{v}" }.join(" ")
        md << "| #{s[:label]} | #{s[:total_calls]} | #{s[:bytes_read]} | #{counts} |\n"
      end
      md << "\n## Candidate missed delegations (runs of >=#{@run_threshold} same-tool calls, distinct targets, no Agent spawn between)\n\n"
      if d[:missed_delegation_candidates].empty?
        md << "None flagged by heuristic.\n"
      else
        d[:missed_delegation_candidates].each do |m|
          md << "- **#{m[:tool]} x#{m[:length]}** (jsonl lines #{m[:lines].join('-')}, #{m[:distinct_targets]} distinct targets)\n"
          m[:targets_sample].each { |t| md << "  - `#{t}`\n" }
        end
      end
      md << "\nHeuristic candidates require auditor judgment; sequential-dependent work is a valid reason not to delegate.\n"
      md
    end

    private

    def build
      lines = parse_lines
      main = lines.reject(&:sidechain?)
      calls, agent_spawns = collect_calls(main)
      attach_results(main, calls)
      attach_agent_metadata(main, calls, agent_spawns)

      orchestrator_calls = calls.values.reject { |c| c.name == "Agent" }
      {
        metadata: metadata(main),
        agent_spawns: agent_spawns.map(&:to_h),
        segments: segments(main).map(&:to_h),
        missed_delegation_candidates:
          MissedDelegationScan.new(orchestrator_calls, @run_threshold).candidates
      }
    end

    def parse_lines
      out = []
      File.foreach(@path).with_index(1) do |raw, i|
        parsed = JSON.parse(raw) rescue next
        out << Line.new(i, parsed)
      end
      out
    end

    def collect_calls(main)
      calls = {}
      spawns = []
      main.each do |line|
        line.tool_uses.each do |block|
          call = ToolCall.new(line, block)
          calls[call.id] = call
          spawns << AgentSpawn.new(call, spawns.size + 1) if call.name == "Agent"
        end
      end
      [calls, spawns]
    end

    def attach_results(main, calls)
      main.each do |line|
        line.tool_results.each do |block|
          call = calls[block["tool_use_id"]] or next
          call.result_chars = result_size(block, line)
        end
      end
    end

    # Async Agent spawns: launch stub carries agentId + resolvedModel in
    # toolUseResult; the real result arrives later as a <task-notification>.
    def attach_agent_metadata(main, calls, spawns)
      by_call_id = spawns.to_h { |s| [s.call.id, s] }
      by_agent_id = {}
      main.each do |line|
        line.tool_results.each do |block|
          spawn = by_call_id[block["tool_use_id"]] or next
          meta = line.raw["toolUseResult"]
          next unless meta.is_a?(Hash)
          spawn.agent_id = meta["agentId"]
          spawn.resolved_model = meta["resolvedModel"]
          by_agent_id[spawn.agent_id] = spawn if spawn.agent_id
        end
      end
      main.each do |line|
        note = line.task_notification or next
        spawn = by_agent_id[note[:task_id]] or next
        spawn.notification_chars = [spawn.notification_chars, note[:chars]].max
      end
    end

    def result_size(block, line)
      content = block["content"]
      size = case content
             when String then content.length
             when Array then content.sum { |b| b.is_a?(Hash) ? b["text"].to_s.length : b.to_s.length }
             else content.to_s.length
             end
      size.zero? ? line.raw["toolUseResult"].to_s.length : size
    end

    def metadata(main)
      assistants = main.select { |l| l.type == "assistant" }
      stamps = main.map(&:timestamp).compact
      {
        cwd: main.map(&:cwd).compact.first,
        started: stamps.min,
        ended: stamps.max,
        assistant_turns: assistants.size,
        human_prompts: main.count(&:human_prompt?),
        main_loop_models: assistants.map(&:model).compact.tally,
        jsonl_lines: main.last&.number
      }
    end

    def segments(main)
      segs = [Segment.new("pre-spawn-1")]
      spawn_count = 0
      main.each do |line|
        line.tool_uses.each do |block|
          call = ToolCall.new(line, block)
          if call.name == "Agent"
            spawn_count += 1
            segs << Segment.new("after-spawn-#{spawn_count}")
          else
            segs.last.add(call)
          end
        end
      end
      resize_segments(segs, main)
    end

    # Re-attach result sizes to the segment copies (segments built fresh above).
    def resize_segments(segs, main)
      sizes = {}
      main.each do |line|
        line.tool_results.each { |b| sizes[b["tool_use_id"]] = result_size(b, line) }
      end
      segs.each do |seg|
        seg.instance_variable_get(:@calls).each { |c| c.result_chars = sizes[c.id].to_i }
      end
      segs
    end
  end
end

if $PROGRAM_NAME == __FILE__
  options = { json: false, run_threshold: 4 }
  OptionParser.new do |o|
    o.banner = "Usage: extract [--json] [--run-threshold N] SESSION.jsonl [...]"
    o.on("--json") { options[:json] = true }
    o.on("--run-threshold N", Integer) { |n| options[:run_threshold] = n }
  end.parse!

  abort "No transcript paths given" if ARGV.empty?
  ARGV.each do |path|
    sheet = OrchAudit::FactSheet.new(path, run_threshold: options[:run_threshold])
    puts options[:json] ? sheet.to_json_s : sheet.to_markdown
    puts
  end
end
