defmodule Refuge.RemoteCache.ByteStream do
  @moduledoc """
  Implementation of the ByteStream service for Bazel Remote Execution API.
  
  The ByteStream service is used for both reading and writing large blobs
  of data, which can be larger than the size supported by a single gRPC message.
  """
  
  use GRPC.Server, service: Google.Bytestream.ByteStream.Service
  
  require Logger
  alias Refuge.RemoteCache.DiskStore
  alias Google.Bytestream.{
    ReadRequest,
    ReadResponse,
    WriteRequest,
    WriteResponse
  }

  @doc """
  Read a blob from the ByteStream service.
  This handles large blob reads by streaming the response.
  """
  @spec read(ReadRequest.t(), GRPC.Server.Stream.t()) :: ReadResponse.t()
  def read(request, stream) do
    Logger.debug("ByteStream.Read request for #{request.resource_name}")
    
    # Parse the resource name to extract the hash
    # Resource names are in the format: "blobs/{hash}/size"
    hash = extract_hash_from_resource(request.resource_name)
    
    # Get the blob data
    case DiskStore.get_cas_blob(hash) do
      {:ok, data} ->
        # If read_offset or read_limit is specified, handle them
        offset = request.read_offset || 0
        limit = request.read_limit || byte_size(data)
        
        # Ensure the offset is valid
        if offset >= byte_size(data) do
          raise GRPC.RPCError, status: :out_of_range, message: "Read offset out of range"
        end
        
        # Get the data chunk
        end_offset = min(offset + limit, byte_size(data))
        chunk = binary_part(data, offset, end_offset - offset)
        
        # Stream the response back
        GRPC.Server.send_reply(stream, %ReadResponse{data: chunk})
        
      {:error, :not_found} ->
        raise GRPC.RPCError, status: :not_found, message: "Resource not found"
      
      {:error, reason} ->
        Logger.error("Failed to read blob: #{inspect(reason)}")
        raise GRPC.RPCError, status: :internal, message: "Internal error reading blob"
    end
  end

  @doc """
  Write a blob to the ByteStream service.
  This handles large blob writes by streaming the request.
  """
  @spec write(Enumerable.t(), GRPC.Server.Stream.t()) :: WriteResponse.t()
  def write(request_enum, _stream) do
    # Accumulate all the chunks first
    {resource_name, data} = 
      Enum.reduce(request_enum, {nil, <<>>}, fn request, {resource_name, acc_data} ->
        # Use the first resource name we see
        resource_name = resource_name || request.resource_name
        
        # Check that resource name is consistent
        if resource_name != request.resource_name do
          raise GRPC.RPCError, 
            status: :invalid_argument, 
            message: "Inconsistent resource name in write request"
        end
        
        # Accumulate the data
        {resource_name, acc_data <> request.data}
      end)
    
    if resource_name == nil do
      raise GRPC.RPCError, status: :invalid_argument, message: "No resource name provided"
    end
    
    # Parse the resource name to extract the hash
    hash = extract_hash_from_resource(resource_name)
    
    # Store the complete blob
    case DiskStore.store_cas_blob(hash, data) do
      :ok ->
        %WriteResponse{committed_size: byte_size(data)}
      
      {:error, reason} ->
        Logger.error("Failed to store blob: #{inspect(reason)}")
        raise GRPC.RPCError, status: :internal, message: "Failed to store blob"
    end
  end

  # Helper functions
  
  defp extract_hash_from_resource(resource_name) do
    # Parse resource names like "blobs/{hash}/size"
    case Regex.run(~r/blobs\/([a-f0-9]{64})/, resource_name) do
      [_, hash] -> hash
      _ ->
        raise GRPC.RPCError, 
          status: :invalid_argument, 
          message: "Invalid resource name: #{resource_name}"
    end
  end
end
