class Multiwoven::Integrations::Source::AwsBedrockModel::Client

def check_connection(connection_config)
  connection_config = connection_config.with_indifferent_access
  create_connection(connection_config)
  model = connection_config[:inference_profile] || connection_config[:model_id]
  payload = format_request(model, connection_config[:request_format])
  @client_runtime.invoke_model(
    model_id: model,
    content_type: "application/json",
    accept: "application/json",
    body: payload
  )
  success_status
rescue StandardError => e
  ConnectionStatus.new(status: ConnectionStatusType["failed"], message: e.message).to_multiwoven_message
end

def create_connection(connection_config)
  @client_runtime = Aws::BedrockRuntime::Client.new(
    region: connection_config[:region],
    access_key_id: connection_config[:access_key],
    secret_access_key: connection_config[:secret_access_key]
  )
end

def discover(_connection_config)
  catalog_json = read_json(CATALOG_SPEC_PATH)
  catalog = build_catalog(catalog_json)
  catalog.to_multiwoven_message
rescue StandardError => e
  handle_exception(e, {
                     context: "AWS:BEDROCK MODEL:DISCOVER:EXCEPTION",
                     type: "error"
                   })
end

def format_request(model, payload)
  case model
  when *MISTRAL_AI_MODEL
    payload_request = JSON.parse(payload)
    prompt = payload_request["prompt"]
    payload_request["prompt"] = "<s>[INST] #{prompt} [/INST]" unless prompt.start_with?("<s>[INST]") && prompt.end_with?("[/INST]")
    payload_request.to_json
  else
    payload
  end
end

def process_response(response)
  data = JSON.parse(response.body.read)
  [RecordMessage.new(data: data, emitted_at: Time.now.to_i).to_multiwoven_message]
end

def read(sync_config)
  connection_config = sync_config.source.connection_specification
  connection_config = connection_config.with_indifferent_access
  payload = sync_config.model.query
  create_connection(connection_config)
  run_model(connection_config, payload)
rescue StandardError => e
  handle_exception(e, {
                     context: "AWS:BEDROCK MODEL:READ:EXCEPTION",
                     type: "error",
                     sync_id: sync_config.sync_id,
                     sync_run_id: sync_config.sync_run_id
                   })
end

def run_model(connection_config, payload)
  model = connection_config[:inference_profile] || connection_config[:model_id]
  payload = format_request(model, payload)
  response = @client_runtime.invoke_model(
    model_id: model,
    content_type: "application/json",
    accept: "application/json",
    body: payload
  )
  process_response(response)
rescue StandardError => e
  handle_exception(e, context: "AWS:BEDROCK MODEL:RUN_MODEL:EXCEPTION", type: "error")
end