defmodule Refuge.RemoteCache.ContentAddressableStorage do
  @moduledoc """
  Implementation of the Bazel Remote Execution ContentAddressableStorage (CAS) service.
  """
  
  use GRPC.Server, service: Build.Bazel.Remote.Execution.V2.ContentAddressableStorage.Service
  
  require Logger
  alias Refuge.RemoteCache.DiskStore
  alias Build.Bazel.Remote.Execution.V2.{
    FindMissingBlobsRequest,
    FindMissingBlobsResponse,
    BatchUpdateBlobsRequest,
    BatchUpdateBlobsResponse,
    BatchReadBlobsRequest,
    BatchReadBlobsResponse,
    BatchReadBlobsResponse.Response,
    GetTreeRequest,
    GetTreeResponse,
    Digest
  }

  @doc """
  Find blobs that are missing from the CAS.
  """
  @spec find_missing_blobs(FindMissingBlobsRequest.t(), GRPC.Server.Stream.t()) :: 
        FindMissingBlobsResponse.t()
  def find_missing_blobs(request, _stream) do
    Logger.debug("FindMissingBlobs request for #{length(request.blob_digests)} blobs")

    # First validate the digests
    invalid_digest = Enum.find(request.blob_digests, fn digest ->
      digest == nil || 
      digest.hash == nil || 
      String.length(digest.hash) != 64 || # Expect SHA-256 = 64 hex chars
      digest.size_bytes < 0
    end)

    if invalid_digest do
      raise GRPC.RPCError, status: :invalid_argument, message: "Invalid digest"
    end

    # Extract hashes from digests
    hashes = Enum.map(request.blob_digests, & &1.hash)
    
    # Find which blobs are missing from the cache
    missing_hashes = DiskStore.find_missing_cas_blobs(hashes)
    
    # Convert missing hashes back to digest objects
    missing_digests = request.blob_digests
      |> Enum.filter(fn digest -> Enum.member?(missing_hashes, digest.hash) end)
    
    Logger.debug("Found #{length(missing_digests)} missing blobs")
    %FindMissingBlobsResponse{missing_blob_digests: missing_digests}
  end

  @doc """
  Update multiple blobs in a batch.
  """
  @spec batch_update_blobs(BatchUpdateBlobsRequest.t(), GRPC.Server.Stream.t()) :: 
        BatchUpdateBlobsResponse.t()
  def batch_update_blobs(request, _stream) do
    Logger.debug("BatchUpdateBlobs request for #{length(request.requests)} blobs")
    
    responses = Enum.map(request.requests, fn req ->
      # Validate the digest
      if req.digest == nil || String.length(req.digest.hash) != 64 do
        raise GRPC.RPCError, status: :invalid_argument, message: "Invalid digest"
      end
      
      # Verify that the data size matches the digest size
      if byte_size(req.data) != req.digest.size_bytes do
        raise GRPC.RPCError, 
          status: :invalid_argument, 
          message: "Data size doesn't match digest size"
      end
      
      # Store the blob
      hash = req.digest.hash
      case DiskStore.store_cas_blob(hash, req.data) do
        :ok ->
          %BatchUpdateBlobsResponse.Response{
            digest: req.digest,
            status: %Google.Rpc.Status{code: 0} # OK status
          }
        {:error, _reason} ->
          %BatchUpdateBlobsResponse.Response{
            digest: req.digest,
            status: %Google.Rpc.Status{code: 13} # INTERNAL error
          }
      end
    end)
    
    %BatchUpdateBlobsResponse{responses: responses}
  end

  @doc """
  Read multiple blobs in a batch.
  """
  @spec batch_read_blobs(BatchReadBlobsRequest.t(), GRPC.Server.Stream.t()) :: 
        BatchReadBlobsResponse.t()
  def batch_read_blobs(request, _stream) do
    Logger.debug("BatchReadBlobs request for #{length(request.digests)} blobs")
    
    responses = Enum.map(request.digests, fn digest ->
      # Validate the digest
      if digest == nil || String.length(digest.hash) != 64 do
        raise GRPC.RPCError, status: :invalid_argument, message: "Invalid digest"
      end
      
      # Retrieve the blob
      hash = digest.hash
      case DiskStore.get_cas_blob(hash) do
        {:ok, data} ->
          %Response{
            digest: digest,
            data: data,
            status: %Google.Rpc.Status{code: 0} # OK status
          }
        {:error, :not_found} ->
          %Response{
            digest: digest,
            data: "",
            status: %Google.Rpc.Status{code: 5} # NOT_FOUND
          }
        {:error, _reason} ->
          %Response{
            digest: digest,
            data: "",
            status: %Google.Rpc.Status{code: 13} # INTERNAL error
          }
      end
    end)
    
    %BatchReadBlobsResponse{responses: responses}
  end

  @doc """
  Get a tree of directory objects rooted at the given digest.
  Note: This implementation is simplified and doesn't fully support the directory tree.
  """
  @spec get_tree(GetTreeRequest.t(), GRPC.Server.Stream.t()) :: GetTreeResponse.t()
  def get_tree(request, _stream) do
    Logger.debug("GetTree request for #{request.root_digest.hash}")
    
    # This is a simplified implementation - for a full implementation,
    # we would need to actually traverse the directory structure
    # For now, we'll just return an empty response
    %GetTreeResponse{
      directories: [],
      next_page_token: ""
    }
  end
end
