require "base64"
require "fileutils"
require "json"
require "uri"

CODEMODE_RESERVED_AGENT_TOOL = "__agent__"
CODEMODE_RESERVED_OUTPUT_TOOL = "__output__"
CODEMODE_RESERVED_SCHEMA_TOOL = "__schema__"
CODEMODE_RESERVED_STORE_TOOL = "__store__"
CODEMODE_RESERVED_LOAD_TOOL = "__load__"
CODEMODE_RESERVED_MCP_TOOL = "__mcp__"
CODEMODE_RESERVED_MEMORY_TOOL = "__memory__"
CODEMODE_RESERVED_APPROVAL_TOOL = "__approval__"
CODEMODE_INTERNAL_URL = Regexp.new("\\A([a-z][a-z0-9+.\\-]*)://(.*)\\z", Regexp::IGNORECASE)

def __codemode_status_enabled?
  connection = $__codemode_connection
  return true unless connection.is_a?(Hash)
  connection["statusEvents"] != false
end

def __codemode_emit_status(op, fields = {}, force: false)
  return unless force || __codemode_status_enabled?
  __codemode_emit({ "type" => "status", "event" => { "op" => op }.merge(fields.transform_keys(&:to_s)) })
end

def __codemode_realpath_prefix(path)
  ancestor = path
  tail = []
  until File.exist?(ancestor)
    parent = File.dirname(ancestor)
    break if parent == ancestor
    tail.unshift(File.basename(ancestor))
    ancestor = parent
  end
  resolved = File.realpath(ancestor)
  tail.empty? ? resolved : File.join(resolved, *tail)
end

def __codemode_resolve_path(value)
  raw = value.to_s
  match = CODEMODE_INTERNAL_URL.match(raw)
  return File.expand_path(raw) unless match

  scheme = match[1].downcase
  roots = $__codemode_connection.is_a?(Hash) ? $__codemode_connection["localRoots"] : nil
  root = roots[scheme] if roots.is_a?(Hash)
  raise "Protocol paths are not supported by this helper: #{raw}" unless root.is_a?(String) && !root.empty?

  relative = URI::DEFAULT_PARSER.unescape(match[2].tr("\\", "/"))
  root_path = File.expand_path(root)
  return root_path if relative.empty?
  if relative.start_with?("/") || relative.split("/").include?("..")
    raise "Unsafe #{scheme}:// path (absolute or traversal): #{raw}"
  end

  resolved = File.expand_path(relative, root_path)
  unless resolved == root_path || resolved.start_with?(root_path + File::SEPARATOR)
    raise "#{scheme}:// path escapes its root: #{raw}"
  end

  real_root = __codemode_realpath_prefix(root_path)
  real_resolved = __codemode_realpath_prefix(resolved)
  unless real_resolved == real_root || real_resolved.start_with?(real_root + File::SEPARATOR)
    raise "#{scheme}:// path escapes its root: #{raw}"
  end
  resolved
end

def __codemode_display_payload(value)
  if value.is_a?(Hash)
    kind = value["type"] || value[:type]
    text_value = value["text"] || value[:text]
    return ["text/markdown", text_value.to_s] if kind == "markdown" && !text_value.nil?
    return ["image/png", value["data"].to_s] if kind == "image" && value["mimeType"] == "image/png"
    return ["image/jpeg", value["data"].to_s] if kind == "image" && value["mimeType"] == "image/jpeg"
    return ["application/json", JSON.generate(value)]
  end
  return ["application/json", JSON.generate(value)] if value.is_a?(Array)
  ["text/plain", value.to_s]
end

def display(value)
  mime_type, payload = __codemode_display_payload(value)
  data_base64 = mime_type.start_with?("image/") ? payload : Base64.strict_encode64(payload)
  __codemode_emit({ "type" => "display", "mimeType" => mime_type, "dataBase64" => data_base64 })
  nil
end

def display_image(base64, mime_type: "image/png")
  __codemode_emit({ "type" => "display", "mimeType" => mime_type.to_s, "dataBase64" => base64.to_s })
  nil
end

def text(data)
  __codemode_emit({ "type" => "text", "stream" => "stdout", "data" => data.to_s })
  nil
end

def print(*values)
  text(values.join)
end

def read(path, offset = 1, limit = nil)
  resolved = __codemode_resolve_path(path)
  data = File.read(resolved, encoding: Encoding::UTF_8)
  if offset > 1 || !limit.nil?
    lines = data.lines
    start = [offset.to_i - 1, 0].max
    data = lines[start, limit || lines.length].to_a.join
  end
  __codemode_emit_status("read", { "path" => resolved, "chars" => data.length })
  data
end

def write(path, content)
  resolved = __codemode_resolve_path(path)
  FileUtils.mkdir_p(File.dirname(resolved))
  data = content.to_s
  File.write(resolved, data)
  __codemode_emit_status("write", { "path" => resolved, "chars" => data.length })
  resolved
end

def env(key = nil, value = nil)
  if key.nil?
    entries = ENV.to_h.sort.to_h
    __codemode_emit_status("env", { "count" => entries.length, "keys" => entries.keys.first(20) })
    return entries
  end

  name = key.to_s
  if value.nil?
    resolved = ENV[name]
    __codemode_emit_status("env", { "key" => name, "action" => "get" })
    return resolved
  end

  resolved = value.to_s
  ENV[name] = resolved
  __codemode_emit_status("env", { "key" => name, "action" => "set" })
  resolved
end

def __codemode_bridge_request(path, payload)
  connection = $__codemode_connection
  raise "Ruby tool bridge is not initialized" unless connection.is_a?(Hash)
  port = connection["port"]
  token = connection["token"]
  raise "Ruby tool bridge is not initialized" unless port.is_a?(Integer) && token.is_a?(String)
  timeout = connection["bridgeTimeoutSeconds"]
  timeout = 120 unless timeout.is_a?(Numeric) && timeout.positive?

  uri = URI("http://127.0.0.1:#{port}#{path}")
  request = Net::HTTP::Post.new(uri)
  request["authorization"] = "Bearer #{token}"
  request["content-type"] = "application/json"
  request.body = JSON.generate(payload)
  __codemode_emit_status("timeout-pause", force: true)
  begin
    response = Net::HTTP.start(uri.hostname, uri.port, open_timeout: 10, read_timeout: timeout) { |http| http.request(request) }
  rescue Net::ReadTimeout
    tool_name = payload["toolName"] || path
    raise "tool call \"#{tool_name}\" timed out after #{timeout}s"
  ensure
    __codemode_emit_status("timeout-resume", force: true)
  end
  body = JSON.parse(response.body.to_s)
  return body["value"] if body.is_a?(Hash) && body["ok"] == true

  error = body.is_a?(Hash) ? body["error"] : body
  raise(error.is_a?(Hash) ? error["message"].to_s : error.to_s)
end

def __codemode_call_tool(name, args)
  __codemode_bridge_request("/call", { "callId" => "rb-#{Process.pid}-#{rand(1_000_000)}", "language" => "rb", "toolName" => name, "args" => args })
end

class CodemodeToolCallable
  def initialize(name)
    @name = name
  end

  def call(args = nil, **kwargs)
    merged = args.nil? ? {} : args.is_a?(Hash) ? args.transform_keys(&:to_s) : raise(ArgumentError, "tool.#{@name}(...) expects a Hash of arguments")
    kwargs.each { |key, value| merged[key.to_s] = value }
    __codemode_call_tool(@name, merged)
  end
end

class CodemodeToolProxy < BasicObject
  def method_missing(name, args = nil, **kwargs)
    ::CodemodeToolCallable.new(name.to_s).call(args, **kwargs)
  end

  def [](name)
    ::CodemodeToolCallable.new(name.to_s)
  end

  def respond_to_missing?(_name, _private = false)
    true
  end
end

def tool
  $__codemode_tool_proxy ||= CodemodeToolProxy.new
end

def completion(prompt, model: "default", system: nil, schema: nil, **kwargs)
  options = { "model" => model }.merge(kwargs.transform_keys(&:to_s))
  options["system"] = system unless system.nil?
  options["schema"] = schema unless schema.nil?
  result = __codemode_bridge_request("/completion", { "prompt" => prompt.to_s, "opts" => options })
  return result unless result.is_a?(Hash)
  return result["value"] if result.key?("value")

  result.fetch("text", result)
end

def tool_schema(name = nil)
  args = name.nil? ? {} : { "name" => name.to_s }
  __codemode_call_tool(CODEMODE_RESERVED_SCHEMA_TOOL, args)
end

def store(key, value)
  raise TypeError, "store() key must be a string" unless key.is_a?(String)
  begin
    JSON.generate(value)
  rescue JSON::GeneratorError
    raise TypeError, "Unable to store \"#{key}\". Only plain serializable values can be stored."
  end
  __codemode_call_tool(CODEMODE_RESERVED_STORE_TOOL, { "language" => "rb", "key" => key, "value" => value })
end

def load(key)
  raise TypeError, "load() key must be a string" unless key.is_a?(String)
  __codemode_call_tool(CODEMODE_RESERVED_LOAD_TOOL, { "language" => "rb", "key" => key })
end

def recall(query, limit = nil)
  raise TypeError, "recall() query must be a non-empty string" unless query.is_a?(String) && !query.empty?
  payload = { "op" => "recall", "query" => query }
  unless limit.nil?
    raise TypeError, "recall() limit must be a positive integer" unless limit.is_a?(Integer) && limit.positive?
    payload["limit"] = limit
  end
  __codemode_call_tool(CODEMODE_RESERVED_MEMORY_TOOL, payload)
end

def approve(tool_name, args = nil, risk = "write")
  raise TypeError, "approve() toolName must be a non-empty string" unless tool_name.is_a?(String) && !tool_name.empty?
  unless %w[read write execute].include?(risk)
    raise TypeError, "approve() risk must be 'read', 'write', or 'execute'"
  end
  payload = { "op" => "check", "toolName" => tool_name, "risk" => risk }
  payload["args"] = args unless args.nil?
  __codemode_call_tool(CODEMODE_RESERVED_APPROVAL_TOOL, payload)
end

class CodemodeMcpHelper
  def list
    __codemode_call_tool(CODEMODE_RESERVED_MCP_TOOL, { "op" => "list" })
  end

  def call(server, tool, args = {})
    raise TypeError, "mcp.call() server must be a non-empty string" unless server.is_a?(String) && !server.empty?
    raise TypeError, "mcp.call() tool must be a non-empty string" unless tool.is_a?(String) && !tool.empty?
    payload = { "op" => "call", "server" => server, "tool" => tool }
    payload["args"] = args.transform_keys(&:to_s) if args.is_a?(Hash)
    __codemode_call_tool(CODEMODE_RESERVED_MCP_TOOL, payload)
  end
end

def mcp
  $__codemode_mcp_helper ||= CodemodeMcpHelper.new
end

def output(*ids, format: "raw", offset: nil, limit: nil)
  raise ArgumentError, "At least one output ID is required" if ids.empty?
  raise ArgumentError, "output() format must be 'raw' or 'tail'" unless ["raw", "tail"].include?(format)
  args = { "ids" => ids.map(&:to_s), "format" => format }
  args["offset"] = offset unless offset.nil?
  args["limit"] = limit unless limit.nil?
  __codemode_call_tool(CODEMODE_RESERVED_OUTPUT_TOOL, args)
end

def agent(prompt, agent: "task", model: nil, label: nil, schema: nil, isolated: nil, apply: nil, merge: nil, handle: false)
  args = { "prompt" => prompt.to_s, "agent" => agent }
  { "model" => model, "label" => label, "schema" => schema, "isolated" => isolated, "apply" => apply, "merge" => merge }.each do |key, value|
    args[key] = value unless value.nil?
  end
  args["handle"] = true if handle
  response = __codemode_call_tool(CODEMODE_RESERVED_AGENT_TOOL, args)
  record = response.is_a?(Hash) ? response : {}
  text_value = record.fetch("text", response)
  result = schema.nil? ? text_value : record.key?("data") ? record["data"] : JSON.parse(text_value.to_s)
  return result unless handle
  { "text" => text_value, "output" => text_value, "handle" => record["handle"] || (record["id"] && "agent://#{record["id"]}"), "id" => record["id"], "agent" => record.fetch("agent", agent) }.tap do |node|
    node["data"] = result unless schema.nil?
  end
end

def __codemode_pool_map(items)
  values = items.to_a
  return [] if values.empty?
  connection = $__codemode_connection
  configured_width = connection.is_a?(Hash) ? connection["parallelPoolWidth"] : nil
  width = configured_width.is_a?(Numeric) ? configured_width.to_i : 4
  workers = [[width, 1].max, values.length].min
  results = Array.new(values.length)
  failures = {}
  failure_mutex = Mutex.new
  queue = Queue.new
  values.each_index { |index| queue << index }
  threads = workers.times.map do
    Thread.new do
      loop do
        index = queue.pop(true) rescue nil
        break if index.nil?
        begin
          results[index] = yield(values[index])
        rescue Exception => error
          failure_mutex.synchronize { failures[index] = error }
        end
      end
    end
  end
  threads.each(&:join)
  raise failures[failures.keys.min] unless failures.empty?
  results
end

def parallel(thunks)
  values = thunks.to_a
  values.each do |thunk|
    raise TypeError, "parallel() expects an iterable of zero-arg callables" unless thunk.respond_to?(:call)
  end
  __codemode_pool_map(values) { |thunk| thunk.call }
end

def pipeline(items, *stages)
  values = items.to_a
  stages.each do |stage|
    raise TypeError, "pipeline() stages must be callables" unless stage.respond_to?(:call)
    values = __codemode_pool_map(values) { |value| stage.call(value) }
  end
  values
end

def log(message)
  __codemode_emit({ "type" => "log", "message" => message.to_s })
  nil
end

def π
  $__codemode_strings || {}
end

def phase(title)
  __codemode_emit({ "type" => "phase", "title" => title.to_s })
  nil
end