#!/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, :rounds

    def initialize(call, index)
      @call = call
      @index = index
      @notification_chars = 0
      @rounds = 1
    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,
        rounds: rounds
      }
    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
      reads = []
      @calls.each do |c|
        counts[c.name] += 1
        next unless READ_TOOLS.include?(c.name)
        bytes += c.result_chars
        reads << { tool: c.name, target: c.target.to_s[0, 200], chars: c.result_chars }
      end
      { label: label, tool_counts: counts, bytes_read: bytes, total_calls: @calls.size, reads: reads }
    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)
      attach_rounds(main, agent_spawns)

      orchestrator_calls = calls.values.reject { |c| c.name == "Agent" }
      {
        metadata: metadata(main, agent_spawns),
        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

    # SendMessage tool_use calls address an agent by its raw agentId or by the
    # teammate name assigned to it; the transcript only ever carries the
    # spawn's description as a name-shaped string, so match on either.
    def attach_rounds(main, spawns)
      targets = send_message_targets(main)
      spawns.each do |spawn|
        spawn.rounds = 1 + targets.count { |t| addressed_to?(t, spawn) }
      end
    end

    def send_message_targets(main)
      main.flat_map do |line|
        line.tool_uses.select { |b| b["name"] == "SendMessage" }
            .map { |b| b.dig("input", "to") }
      end.compact
    end

    def addressed_to?(target, spawn)
      t = target.to_s.strip.downcase
      return false if t.empty?
      return true if spawn.agent_id && t == spawn.agent_id.to_s.downcase
      desc = spawn.call.input["description"].to_s.strip.downcase
      !desc.empty? && (t == desc || desc.include?(t) || t.include?(desc))
    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, spawns)
      assistants = main.select { |l| l.type == "assistant" }
      stamps = main.map(&:timestamp).compact
      main_tok = main_loop_output_tokens(assistants)
      side_tok = sidechain_output_tokens
      rounds = spawns.map(&:rounds)
      {
        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,
        main_loop_output_tokens: main_tok,
        sidechain_output_tokens: side_tok,
        main_loop_share: main_loop_share(main_tok, side_tok),
        total_spawns: spawns.size,
        max_rounds: rounds.max || 0
      }
    end

    def main_loop_output_tokens(assistants)
      assistants.sum { |l| l.message.dig("usage", "output_tokens").to_i }
    end

    # Subagent transcripts persist as sibling files at
    # <dir>/<main-basename-without-.jsonl>/subagents/agent-*.jsonl (same
    # native JSONL shape, own usage blocks); 0 when the dir is absent.
    def sidechain_output_tokens
      dir = File.dirname(@path)
      base = File.basename(@path, ".jsonl")
      total = 0
      Dir.glob(File.join(dir, base, "subagents", "*.jsonl")).each do |f|
        File.foreach(f) do |raw|
          parsed = JSON.parse(raw) rescue next
          next unless parsed["type"] == "assistant"
          total += parsed.dig("message", "usage", "output_tokens").to_i
        end
      end
      total
    end

    def main_loop_share(main_tok, side_tok)
      return 1.0 if side_tok.zero?
      (main_tok.to_f / (main_tok + side_tok)).round(3)
    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