How to set cookies to remember country code for Remix + Hydrogen Custom Storefront

Question

So I have a custom storefront built on Hydrogen + Remix, and this is my server.tsx file.

// Virtual entry point for the app
import * as remixBuild from '@remix-run/dev/server-build';
import {createRequestHandler} from '@remix-run/server-runtime';
import {createStorefrontClient} from '@shopify/hydrogen';
import { getStorefrontHeaders } from '@shopify/remix-oxygen';

import {HydrogenSession} from '~/lib/session.server';
import {getApproximateLocaleFromRequest, getLocaleFromRequest} from '~/lib/utils';

/**
 * Export a fetch handler in module format.
 */
export default async function (request: Request): Promise {
  try {
    /**
     * This has to be done so messy because process.env can't be destructured
     * and only variables explicitly named are present inside a Vercel Edge Function.
     * See https://github.com/vercel/next.js/pull/31237/files
     */
    const env: Env = {
      SESSION_SECRET: '',
      PUBLIC_STOREFRONT_API_TOKEN: '',
      PRIVATE_STOREFRONT_API_TOKEN: '',
      PUBLIC_STORE_DOMAIN: '',
      PUBLIC_STOREFRONT_ID: ''
    };
    env.SESSION_SECRET = process.env.SESSION_SECRET;
    env.PUBLIC_STOREFRONT_API_TOKEN = process.env.PUBLIC_STOREFRONT_API_TOKEN;
    env.PRIVATE_STOREFRONT_API_TOKEN = process.env.PRIVATE_STOREFRONT_API_TOKEN;
    env.PUBLIC_STORE_DOMAIN = process.env.PUBLIC_STORE_DOMAIN;
    env.PUBLIC_STOREFRONT_ID = process.env.PUBLIC_STOREFRONT_ID;
    /**
     * Open a cache instance in the worker and a custom session instance.
     */
    if (!env?.SESSION_SECRET) {
      throw new Error('SESSION_SECRET process.environment variable is not set');
    }

    const [session] = await Promise.all([
      HydrogenSession.init(request, [process.env.SESSION_SECRET]),
    ]);

    /**
     * Create Hydrogen's Storefront client.
     */
    const {storefront} = createStorefrontClient({
      i18n: getLocaleFromRequest(request),
      publicStorefrontToken: env.PUBLIC_STOREFRONT_API_TOKEN,
      privateStorefrontToken: env.PRIVATE_STOREFRONT_API_TOKEN,
      storeDomain: `https://${env.PUBLIC_STORE_DOMAIN}`,
      storefrontApiVersion: env.PUBLIC_STOREFRONT_API_VERSION || '2023-04',
      storefrontId: env.PUBLIC_STOREFRONT_ID,
      storefrontHeaders: getStorefrontHeaders(request),
    });

    const handleRequest = createRequestHandler(remixBuild as any, 'production');

    const response = await handleRequest(request, {
      session,
      storefront,
      env,
      waitUntil: () => Promise.resolve(),
    });

    if (response.status === 404) {
      /**
       * Check for redirects only when there's a 404 from the app.
       * If the redirect doesn't exist, then `storefrontRedirect`
       * will pass through the 404 response.
       */
      // return storefrontRedirect({request, response, storefront});
    }

    return response;
  } catch (error) {
    // eslint-disable-next-line no-console
    console.error(error);
    return new Response('An unexpected error occurred', {status: 500});
  }
}

The current implementation for getLocaleFromRequest is

import {countries} from './countries';

export function getLocaleFromRequest(request: Request): I18nLocale {
  const url = new URL(request.url);
  const firstPathPart =
    '/' + url.pathname.substring(1).split('/')[0].toLowerCase();

  return countries[firstPathPart]
    ? {
        ...countries[firstPathPart],
        pathPrefix: firstPathPart,
      }
    : {
        ...countries['/gb'],
        pathPrefix: '',
      };
}

The countries.tsx import contains the following:

import type {Localizations} from '~/lib/type';

export const countries: Localizations = {
  '/us': {
    label: 'United States (USD $)',
    language: 'EN',
    country: 'US',
    currency: 'USD',
  },
  '/au': {
    label: 'Australia (AUD $)',
    language: 'EN',
    country: 'AU',
    currency: 'AUD',
  },
  '/ca': {
    label: 'Canada (CAD $)',
    language: 'EN',
    country: 'CA',
    currency: 'CAD',
  },
  '/fr': {
    label: 'France (EUR €)',
    language: 'EN',
    country: 'FR',
    currency: 'EUR',
  },
  '/gb': {
    label: 'United Kingdom (GBP £)',
    language: 'EN',
    country: 'GB',
    currency: 'GBP',
  },
};

I tried an implementation where I try to use cookies to remember the country code so that, if the user doesn't type in the country code in the URL, it remembers the country code from the cookie instead of defaulting to GB.

export async function getApproximateLocaleFromRequest(
  request: Request,
): I18nLocale {
  const cookieHeader = request.headers.get("Cookie");
  const cookie = await getCountryCode.parse(cookieHeader);

  const url = new URL(request.url);
  const firstPathPart =
    '/' + url.pathname.substring(1).split('/')[0].toLowerCase();

  if (countries[firstPathPart]) {
    // Set cookie
    cookie.country_code = countries[firstPathPart].country.toLocaleLowerCase();
    request.headers.set("Cookie", await getCountryCode.serialize({
      cookie
    }));

    return {
      ...countries[firstPathPart],
      pathPrefix: firstPathPart,
    };
  } else {
    if (!cookie?.country_code) {
      // If there's no locale cookie, it means it's the user's first time visiting.
      return {
        ...countries['/fr'],
        pathPrefix: '',
      };
    
    } else {
      return {
        ...countries[`/${cookie.country_code}`],
        pathPrefix: cookie.country_code,
      };
    }
  }
}

This implementation is incorrect, and I need guidance on how to achieve this. Any help would be much appreciated.