Session Key Management

"use client";

import { useAccount } from "wagmi";
import { useAbstractClient } from "@abstract-foundation/agw-react";
import { Button } from "@/components/ui/button";
import {
  DropdownMenuItem,
  DropdownMenuSeparator,
} from "@/components/ui/dropdown-menu";
import { ConnectWalletButton } from "@/components/connect-wallet-button";
import { useSessionKey } from "@/hooks/use-session-key";
import { useCreateSessionKey } from "@/hooks/use-create-session-key";
import { useRevokeSessionKey } from "@/hooks/use-revoke-session-key";
import { cn } from "@/lib/utils";
import { type ClassValue } from "clsx";

interface SessionKeyButtonProps {
  className?: ClassValue;
}

/**
 * Session Key Button
 * - Not connected: Shows "Connect Wallet" button
 * - Connected but no session: Shows "Create Session Key" button
 * - Session exists: Shows "Session Active" with dropdown containing revoke option
 */
export function SessionKeyButton({ className }: SessionKeyButtonProps) {
  // Wallet state
  const { isConnected, isConnecting, isReconnecting } = useAccount();
  const { data: abstractClient, isLoading: isAbstractClientLoading } =
    useAbstractClient();

  // Session key state
  const {
    data: sessionData,
    isLoading: isSessionLoading,
    isFetched: isSessionFetched,
    isReadyForCreate,
  } = useSessionKey();

  // Create and revoke sesssion keys from the connected wallet
  const createSessionMutation = useCreateSessionKey();
  const { revokeSession, isPending: isRevoking } = useRevokeSessionKey();

  // Check if user has an active sessiona
  const hasActiveSession = !!sessionData;

  const isChecking =
    isSessionLoading ||
    isConnecting ||
    isReconnecting ||
    isAbstractClientLoading ||
    (isConnected && (!abstractClient || !isSessionFetched));

  // Handle session creation
  const handleCreateSession = () => {
    createSessionMutation.mutate();
  };

  // Handle session revocation
  const handleRevokeSession = async (e: React.MouseEvent) => {
    // Prevent dropdown from closing
    e.preventDefault();
    e.stopPropagation();

    if (sessionData?.session) {
      await revokeSession(sessionData.session);
    }
  };

  // Not connected: Show Connect Wallet button
  if (!isConnected) {
    return <ConnectWalletButton className={className} />;
  }

  // Connected and has active session: Show session status with dropdown
  if (isConnected && hasActiveSession && !isChecking) {
    return (
      <ConnectWalletButton
        className={className}
        customDropdownItems={[
          <DropdownMenuSeparator key="sep" />,
          <DropdownMenuItem
            key="session-info"
            className="focus:bg-transparent cursor-auto"
          >
            <div className="flex items-center justify-between w-full">
              <span className="text-xs text-muted-foreground">Session:</span>
              <div className="flex items-center space-x-2">
                <div className="h-2 w-2 bg-green-500 rounded-full" />
                <span className="text-xs text-green-600 dark:text-green-400">
                  Active
                </span>
              </div>
            </div>
          </DropdownMenuItem>,
          <DropdownMenuSeparator key="sep2" />,
          <DropdownMenuItem
            key="revoke"
            onClick={handleRevokeSession}
            disabled={isRevoking}
            className="text-destructive"
          >
            {isRevoking ? (
              <>
                <Spinner className="mr-2 h-4 w-4 animate-spin" />
                Revoking...
              </>
            ) : (
              <>
                <RevokeIcon className="mr-2 h-4 w-4" />
                Revoke Session Key
              </>
            )}
          </DropdownMenuItem>,
        ]}
      />
    );
  }

  // If we're connected but not yet ready for create (first fetch pending), show Checking
  if (isConnected && !hasActiveSession && !isReadyForCreate) {
    return (
      <Button
        disabled
        className={cn("cursor-pointer group min-w-40", className)}
      >
        <Spinner className="mr-2 h-4 w-4 animate-spin" />
        Checking...
      </Button>
    );
  }

  // Connected but no session OR loading: Show Create Session Key button
  return (
    <Button
      onClick={handleCreateSession}
      disabled={createSessionMutation.isPending || isChecking}
      className={cn("cursor-pointer group min-w-40", className)}
    >
      {createSessionMutation.isPending ? (
        <>
          <Spinner className="mr-2 h-4 w-4 animate-spin" />
          Creating...
        </>
      ) : isChecking ? (
        <>
          <Spinner className="mr-2 h-4 w-4 animate-spin" />
          Checking...
        </>
      ) : (
        <>
          <KeyIcon className="mr-2 h-4 w-4" />
          Create Session Key
        </>
      )}
    </Button>
  );
}

function Spinner({ className }: { className?: ClassValue }) {
  return (
    <svg
      className={cn("animate-spin", className)}
      xmlns="http://www.w3.org/2000/svg"
      fill="none"
      viewBox="0 0 24 24"
    >
      <circle
        className="opacity-25"
        cx="12"
        cy="12"
        r="10"
        stroke="currentColor"
        strokeWidth="4"
      />
      <path
        className="opacity-75"
        fill="currentColor"
        d="m4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z"
      />
    </svg>
  );
}

function KeyIcon({ className }: { className?: ClassValue }) {
  return (
    <svg
      className={cn(className)}
      xmlns="http://www.w3.org/2000/svg"
      fill="none"
      viewBox="0 0 24 24"
      stroke="currentColor"
      strokeWidth="2"
    >
      <path
        strokeLinecap="round"
        strokeLinejoin="round"
        d="M15.75 5.25a3 3 0 013 3m3 0a6 6 0 01-7.029 5.912c-.563-.097-1.159.026-1.563.43L10.5 17.25H8.25v2.25H6v2.25H2.25v-2.818c0-.597.237-1.17.659-1.591l6.499-6.499c.404-.404.527-1 .43-1.563A6 6 0 1121.75 8.25z"
      />
    </svg>
  );
}

function RevokeIcon({ className }: { className?: ClassValue }) {
  return (
    <svg
      className={cn(className)}
      xmlns="http://www.w3.org/2000/svg"
      fill="none"
      viewBox="0 0 24 24"
      stroke="currentColor"
      strokeWidth="2"
    >
      <path
        strokeLinecap="round"
        strokeLinejoin="round"
        d="M6 18L18 6M6 6l12 12"
      />
    </svg>
  );
}

"use client";

import { useQuery } from "@tanstack/react-query";
import { useAccount } from "wagmi";
import { useAbstractClient } from "@abstract-foundation/agw-react";
import { getStoredSession } from "../lib/get-stored-session-key";

/**
 * Hook to retrieve and validate the stored Abstract session
 * @returns The session data with loading and error states
 */
export function useSessionKey() {
  const { address, isConnecting, isReconnecting } = useAccount();
  const { data: abstractClient, isLoading: isAbstractClientLoading } =
    useAbstractClient();

  const queryResult = useQuery({
    queryKey: ["session-key", address],
    queryFn: async () => {
      // These should never happen due to enabled condition, but adding for type safety
      if (!abstractClient) {
        throw new Error("No Abstract client found");
      }
      if (!address) {
        throw new Error("No wallet address found");
      }

      return await getStoredSession(abstractClient, address);
    },
    enabled: !!address && !!abstractClient,
    staleTime: 0, // Always refetch when invalidated
    retry: (failureCount, error) => {
      // Don't retry if it's a client/address error (these won't resolve with retry)
      if (
        error.message.includes("No Abstract client") ||
        error.message.includes("No wallet address")
      ) {
        return false;
      }
      // Retry network/validation errors up to 2 times
      return failureCount < 2;
    },
    retryDelay: (attemptIndex) => Math.min(1000 * 2 ** attemptIndex, 30000), // Exponential backoff
  });

  const finalIsLoading =
    isConnecting ||
    isReconnecting ||
    isAbstractClientLoading ||
    queryResult.isLoading ||
    (!!address && !abstractClient); // Address ready but Abstract client not ready yet

  // Ready to show the "Create Session Key" state only after we have a connected
  // wallet, an Abstract client, and the first query has settled with no session
  const isReadyForCreate =
    !!address &&
    !!abstractClient &&
    !!queryResult.isFetched &&
    !finalIsLoading &&
    !queryResult.data;

  // Debug logging removed

  return {
    ...queryResult,
    isLoading: finalIsLoading,
    isReadyForCreate,
  };
}

"use client";

import { useMutation, useQueryClient } from "@tanstack/react-query";
import { useAccount } from "wagmi";
import { useAbstractClient } from "@abstract-foundation/agw-react";
import { createAndStoreSession } from "../lib/create-and-store-session-key";
import { toast } from "sonner";

/**
 * Hook to create and store Abstract sessions
 * @returns Mutation functions and state for creating sessions
 */
export function useCreateSessionKey() {
    const { data: abstractClient } = useAbstractClient();
    const { address } = useAccount();
    const queryClient = useQueryClient();

    const createSessionMutation = useMutation({
        mutationFn: async () => {
            if (!address) {
                throw new Error("No wallet address found");
            }
            if (!abstractClient) {
                throw new Error("No Abstract client found");
            }

            return createAndStoreSession(abstractClient, address);
        },
        onSuccess: async () => {
            // Invalidate the session query to force a refetch
            await queryClient.invalidateQueries({
                queryKey: ["session-key", address],
            });
            
            // Also refetch immediately to ensure state updates
            await queryClient.refetchQueries({
                queryKey: ["session-key", address],
            });
            
            toast.success("Session key created successfully");
        },
        onError: (error) => {
            console.error("Failed to create session:", error);
            toast.error("Failed to create session key");
        },
    });

    return createSessionMutation;
}

"use client";

import { useRevokeSessions } from "@abstract-foundation/agw-react";
import { useAccount } from "wagmi";
import { useQueryClient } from "@tanstack/react-query";
import { toast } from "sonner";
import { clearStoredSession } from "../lib/clear-stored-session-key";
import type { SessionConfig } from "@abstract-foundation/agw-client/sessions";

/**
 * Hook to revoke session keys with proper cleanup and error handling
 * @returns Mutation functions and state for revoking sessions
 */
export function useRevokeSessionKey() {
    const { address } = useAccount();
    const { revokeSessionsAsync, isPending, isError, error } = useRevokeSessions();
    const queryClient = useQueryClient();

    const revokeSession = async (session: SessionConfig) => {
        if (!address) {
            toast.error("No wallet address found");
            return;
        }

        try {
            const result = await revokeSessionsAsync({
                sessions: session
            });

            // The result might be undefined or have different structure; keep for potential future use

            // Clear local storage after successful revocation
            clearStoredSession(address);
            
            // Invalidate the session query to force a refetch
            await queryClient.invalidateQueries({
                queryKey: ["session-key", address],
            });
            
            // Also refetch immediately to ensure state updates
            await queryClient.refetchQueries({
                queryKey: ["session-key", address],
            });

            toast.success("Session key revoked successfully");
        } catch (err) {
            console.error("Failed to revoke session:", err);
            toast.error("Failed to revoke session key");
            throw err;
        }
    };

    return {
        revokeSession,
        isPending,
        isError,
        error
    };
}

"use client";

import { useMutation, useQueryClient } from "@tanstack/react-query";
import { useAccount } from "wagmi";
import { useAbstractClient } from "@abstract-foundation/agw-react";
import { getStoredSession } from "@/lib/get-stored-session-key";
import { publicClient } from "@/config/viem-clients";
import { chain } from "@/config/chain";
import { privateKeyToAccount } from "viem/accounts";
import {
  Address,
  Abi,
  ContractFunctionName,
  ContractFunctionArgs,
  WriteContractParameters,
} from "viem";

// Generic type-safe transaction parameters for any contract ABI
type SessionKeyTransactionParams<
  TAbi extends Abi,
  TFunctionName extends ContractFunctionName<
    TAbi,
    "payable" | "nonpayable"
  > = ContractFunctionName<TAbi, "payable" | "nonpayable">
> = {
  abi: TAbi;
  address: Address;
  functionName: TFunctionName;
  args?: ContractFunctionArgs<TAbi, "payable" | "nonpayable", TFunctionName>;
  value?: bigint;
};

interface UseSessionKeyTransactionOptions {
  onSuccess?: (data: { hash: Address }) => void;
  onError?: (error: Error) => void;
}

/**
 * Hook for submitting contract transactions using session keys
 * Provides gasless transactions for approved contract functions
 *
 * @example
 * ```typescript
 * const { mutate, mutateAsync } = useSessionKeyTransaction();
 *
 * // Type-safe usage with any contract
 * mutate({
 *   abi: MY_CONTRACT_ABI,
 *   address: MY_CONTRACT_ADDRESS,
 *   functionName: "transfer", // ✅ Autocomplete available based on ABI
 *   args: [recipient, amount], // ✅ TypeScript knows the correct argument types
 *   value: parseEther("0.1")
 * });
 * ```
 */
export function useSessionKeyTransaction(
  options?: UseSessionKeyTransactionOptions
) {
  const { address } = useAccount();
  const { data: abstractClient } = useAbstractClient();
  const queryClient = useQueryClient();

  const transactionMutation = useMutation<
    { hash: Address },
    Error,
    SessionKeyTransactionParams<Abi>
  >({
    mutationFn: async ({
      abi,
      address: contractAddress,
      functionName,
      args,
      value = BigInt(0),
    }) => {
      if (!address) {
        throw new Error("No wallet address found");
      }
      if (!abstractClient) {
        throw new Error("No Abstract client found");
      }

      // Get stored session
      const sessionData = await getStoredSession(abstractClient, address);
      if (!sessionData) {
        throw new Error(
          "No valid session found. Please create a session key first."
        );
      }

      const sessionClient = abstractClient.toSessionClient(
        privateKeyToAccount(sessionData.privateKey),
        sessionData.session
      );

      // Submit transaction using session key
      const txHash = await sessionClient.writeContract({
        abi,
        account: sessionClient.account,
        address: contractAddress,
        functionName,
        args,
        value,
        chain,
      } as WriteContractParameters);

      if (!txHash) {
        throw new Error("Transaction failed - no hash returned");
      }

      // Wait for transaction confirmation
      const receipt = await publicClient.waitForTransactionReceipt({
        hash: txHash,
        timeout: 10000, // 10 second timeout - if it takes longer, likely nonce issue
        pollingInterval: 1000 / 3, // Check 3 times per second
      });

      if (receipt.status !== "success") {
        // Try to get the revert reason from the transaction
        let revertReason = "Unknown revert reason";
        try {
          // Simulate the transaction to get the revert reason
          await publicClient.call({
            to: contractAddress,
            data: (await publicClient.getTransaction({ hash: txHash })).input,
            blockNumber: receipt.blockNumber,
          });
        } catch (simulationError: any) {
          if (simulationError.message) {
            revertReason = simulationError.message;
          }
        }

        throw new Error(
          `Transaction failed during execution: ${receipt.status}. Reason: ${revertReason}`
        );
      }

      return { hash: txHash };
    },
    onSuccess: (data: { hash: Address }) => {
      // Invalidate session key queries to refresh UI state
      queryClient.invalidateQueries({
        queryKey: ["session-key"],
      });
      options?.onSuccess?.(data);
    },
    onError: (error: Error) => {
      console.error("Session key transaction failed:", error);
      options?.onError?.(error);
    },
  });

  // Type-safe generic mutate function
  const mutate = <TAbi extends Abi>(
    params: SessionKeyTransactionParams<TAbi>
  ) => {
    return transactionMutation.mutate(
      params as SessionKeyTransactionParams<Abi>
    );
  };

  const mutateAsync = <TAbi extends Abi>(
    params: SessionKeyTransactionParams<TAbi>
  ) => {
    return transactionMutation.mutateAsync(
      params as SessionKeyTransactionParams<Abi>
    );
  };

  return {
    // Type-safe generic mutation functions
    mutate,
    mutateAsync,
    // Mutation state
    isPending: transactionMutation.isPending,
    isError: transactionMutation.isError,
    isSuccess: transactionMutation.isSuccess,
    error: transactionMutation.error,
    data: transactionMutation.data,
    reset: transactionMutation.reset,
  };
}

/**
 * IMPORTANT: https://docs.abs.xyz/abstract-global-wallet/session-keys/going-to-production
 * Your session key config requires approval to operate on Abstract mainnet via whitelist.
 */

import {
    LimitType,
    type SessionConfig,
} from "@abstract-foundation/agw-client/sessions";
import { parseEther, toFunctionSelector } from "viem";

/**
 * What call policies you wish to allow for the session key
 * Learn more: https://docs.abs.xyz/abstract-global-wallet/agw-client/session-keys/createSession#param-call-policies
 */
export const CALL_POLICIES = [
    {
        target: "0xC4822AbB9F05646A9Ce44EFa6dDcda0Bf45595AA" as `0x${string}`, // Contract address
        selector: toFunctionSelector("mint(address,uint256)"), // Allowed function
        // Gas parameters
        valueLimit: {
            limitType: LimitType.Unlimited,
            limit: BigInt(0),
            period: BigInt(0),
        },
        maxValuePerUse: BigInt(0),
        constraints: [],
    }
];

/**
 * What transfer policies you wish to allow for the session key
 * Learn more: https://docs.abs.xyz/abstract-global-wallet/agw-client/session-keys/createSession#param-transfer-policies
 */
export const TRANSFER_POLICIES = [
    // ... Your transfer policies here
]

export const SESSION_KEY_CONFIG: Omit<SessionConfig, "signer"> = {
    expiresAt: BigInt(Math.floor(Date.now() / 1000) + 60 * 60 * 24 * 30), // 30 days from now
    feeLimit: {
        limitType: LimitType.Lifetime,
        limit: parseEther("1"), // 1 ETH lifetime gas limit
        period: BigInt(0),
    },
    callPolicies: CALL_POLICIES,
    transferPolicies: [],
};

import type { Address, Hex } from "viem";
import type { AbstractClient } from "@abstract-foundation/agw-client";
import { getSessionHash, type SessionConfig } from "@abstract-foundation/agw-client/sessions";
import { LOCAL_STORAGE_KEY_PREFIX, getEncryptionKey, decrypt } from "./session-encryption-utils";
import { validateSession } from "./validate-session-key";
import { CALL_POLICIES } from "@/config/session-key-policies";

export interface StoredSessionData {
    session: SessionConfig;
    privateKey: Hex;
}

/**
 * @function getStoredSession
 * @description Retrieves, decrypts, and validates a stored session for a wallet address
 * 
 * This function performs several steps to securely retrieve and validate a stored session:
 * 1. Checks local storage for encrypted session data under the wallet address key
 * 2. Retrieves the encryption key for the wallet address
 * 3. Decrypts the session data using the encryption key
 * 4. Parses the decrypted data to obtain session information
 * 5. Validates that call policies match current configuration
 * 6. Validates the session by checking its status on-chain
 * 
 * @param {AbstractClient} abstractClient - The Abstract client instance
 * @param {Address} userAddress - The wallet address to retrieve session for
 * @returns {Promise<StoredSessionData | null>} The session data if valid, null otherwise
 */
export const getStoredSession = async (
    abstractClient: AbstractClient,
    userAddress: Address
): Promise<StoredSessionData | null> => {
    if (!userAddress) return null;

    const encryptedData = localStorage.getItem(
        `${LOCAL_STORAGE_KEY_PREFIX}${userAddress}`
    );

    if (!encryptedData) return null;

    try {
        const key = await getEncryptionKey(userAddress);
        const decryptedData = await decrypt(encryptedData, key);

        const sessionData: StoredSessionData = JSON.parse(decryptedData, (_, value) => {
            // Handle bigint deserialization
            if (typeof value === "string" && /^\d+$/.test(value)) {
                try {
                    return BigInt(value);
                } catch {
                    return value;
                }
            }
            return value;
        });

        // Check if stored call policies match current configuration
        const storedPoliciesJson = JSON.stringify(
            sessionData.session.callPolicies,
            (_, value) => typeof value === "bigint" ? value.toString() : value
        );
        const currentPoliciesJson = JSON.stringify(
            CALL_POLICIES,
            (_, value) => typeof value === "bigint" ? value.toString() : value
        );

        if (storedPoliciesJson !== currentPoliciesJson) {
            // Call policies changed; invalidate stored session to force refresh
            return null;
        }

        // Validate the session is still active on-chain
        const sessionHash = getSessionHash(sessionData.session);
        const isValid = await validateSession(
            abstractClient,
            userAddress,
            sessionHash
        );

        return isValid ? sessionData : null;
    } catch (error) {
        console.error("Failed to retrieve stored session:", error);
        return null;
    }
};

import { generatePrivateKey, privateKeyToAccount } from "viem/accounts";
import { Address } from "viem";
import { SessionConfig } from "@abstract-foundation/agw-client/sessions";
import { LOCAL_STORAGE_KEY_PREFIX, getEncryptionKey, encrypt } from "./session-encryption-utils";
import { AbstractClient } from "@abstract-foundation/agw-client";
import { publicClient } from "@/config/viem-clients";
import { SESSION_KEY_CONFIG } from "@/config/session-key-policies";

/**
 * @function createAndStoreSession
 * @description Creates a new Abstract Global Wallet session and stores it securely in local storage
 *
 * @param {Address} userAddress - The wallet address that will own the session
 *
 * @returns {Promise<Object|null>} A promise that resolves to:
 *   - The created session data object (containing `session` and `privateKey`) if successful
 *   - null if the userAddress is empty or invalid
 *
 * @throws {Error} Throws "Session creation failed" if there's an error during session creation
 */
export const createAndStoreSession = async (
    abstractClient: AbstractClient,
    userAddress: Address
): Promise<{
    session: SessionConfig;
    privateKey: Address;
} | null> => {
    if (!userAddress) return null;

    try {
        const sessionPrivateKey = generatePrivateKey();
        const sessionSigner = privateKeyToAccount(sessionPrivateKey);

        const { session, transactionHash } = await abstractClient.createSession({
            session: {
                signer: sessionSigner.address,
                ...SESSION_KEY_CONFIG,
            },
        });

        if (transactionHash) {
            await publicClient.waitForTransactionReceipt({
                hash: transactionHash,
            });
        } else {
            throw new Error("Transaction hash is null. Session was not created.");
        }

        const sessionData = { session, privateKey: sessionPrivateKey };
        const key = await getEncryptionKey(userAddress);
        const encryptedData = await encrypt(
            JSON.stringify(sessionData, (_, value) =>
                typeof value === "bigint" ? value.toString() : value
            ),
            key
        );

        localStorage.setItem(
            `${LOCAL_STORAGE_KEY_PREFIX}${userAddress}`,
            encryptedData
        );
        return sessionData;
    } catch (error) {
        console.error("Failed to create session:", error);
        throw new Error(`Failed to create session key, ${error}`);
    }
};

import { LOCAL_STORAGE_KEY_PREFIX, ENCRYPTION_KEY_PREFIX } from "./session-encryption-utils";
import type { Address } from "viem";

/**
 * @function clearStoredSession
 * @description Removes all stored session data for a specific wallet address from local storage
 *
 * This function cleans up both the encrypted session data and the encryption key
 * associated with a wallet address from the browser's local storage. It's typically
 * used when a session has expired, been revoked, or when the user wants to clear
 * their session data for privacy/security reasons.
 *
 * The function removes two items from local storage:
 * 1. The encrypted session data (stored with LOCAL_STORAGE_KEY_PREFIX + address)
 * 2. The encryption key used to encrypt/decrypt the session (stored with ENCRYPTION_KEY_PREFIX + address)
 *
 * @param {Address} userAddress - The wallet address whose session data should be cleared
 */
export const clearStoredSession = (userAddress: Address) => {
    localStorage.removeItem(`${LOCAL_STORAGE_KEY_PREFIX}${userAddress}`);
    localStorage.removeItem(`${ENCRYPTION_KEY_PREFIX}${userAddress}`);
};

import type { Address } from "viem";


/**
 * @constant {string} LOCAL_STORAGE_KEY_PREFIX
 * @description Prefix used for storing encrypted session data in local storage
 *
 * The actual storage key is created by appending the user's wallet address to this prefix,
 * ensuring each wallet address has its own unique storage key.
 * The prefix includes the current NODE_ENV to separate data between environments.
 */
export const LOCAL_STORAGE_KEY_PREFIX = `abstract_session_${process.env.NODE_ENV || "development"}_`;

/**
 * @constant {string} ENCRYPTION_KEY_PREFIX
 * @description Prefix used for storing encryption keys in local storage
 *
 * The actual storage key is created by appending the user's wallet address to this prefix,
 * ensuring each wallet address has its own unique encryption key stored separately from the
 * encrypted session data.
 * The prefix includes the current NODE_ENV to separate data between environments.
 */
export const ENCRYPTION_KEY_PREFIX = `encryption_key_${process.env.NODE_ENV || "development"}_`;


/**
 * @function getEncryptionKey
 * @description Retrieves or generates an AES-GCM encryption key for a specific wallet address
 *
 * This function manages encryption keys used to secure session data in local storage.
 * It first checks if an encryption key already exists for the given wallet address.
 * If found, it imports and returns the existing key. Otherwise, it generates a new
 * 256-bit AES-GCM key, stores it in local storage, and returns it.
 *
 * The encryption keys are stored in local storage with a prefix (defined in constants.ts)
 * followed by the wallet address to ensure each wallet has its own unique encryption key.
 *
 * @param {Address} userAddress - The wallet address to get or generate an encryption key for
 *
 * @returns {Promise<CryptoKey>} A promise that resolves to a CryptoKey object that can be
 *                              used with the Web Crypto API for encryption and decryption
 */
export const getEncryptionKey = async (
    userAddress: Address
): Promise<CryptoKey> => {
    const storedKey = localStorage.getItem(
        `${ENCRYPTION_KEY_PREFIX}${userAddress}`
    );

    if (storedKey) {
        return crypto.subtle.importKey(
            "raw",
            Buffer.from(storedKey, "hex"),
            { name: "AES-GCM" },
            false,
            ["encrypt", "decrypt"]
        );
    }

    const key = await crypto.subtle.generateKey(
        { name: "AES-GCM", length: 256 },
        true,
        ["encrypt", "decrypt"]
    );

    const exportedKey = await crypto.subtle.exportKey("raw", key);
    localStorage.setItem(
        `${ENCRYPTION_KEY_PREFIX}${userAddress}`,
        Buffer.from(exportedKey).toString("hex")
    );

    return key;
};

/**
 * @function encrypt
 * @description Encrypts data using AES-GCM encryption with a provided CryptoKey
 *
 * This function uses the Web Crypto API to encrypt session data for secure storage
 * in the browser's local storage. It generates a random initialization vector (IV)
 * for each encryption operation to ensure security. The encrypted data and IV are
 * both stored in the returned JSON string.
 *
 * @param {string} data - The data to encrypt, typically a stringified JSON object
 *                        containing session information and private keys
 * @param {CryptoKey} key - The AES-GCM encryption key to use
 *
 * @returns {Promise<string>} A promise that resolves to a JSON string containing
 *                           the encrypted data and the initialization vector (IV)
 *                           both encoded as hex strings
 */
export const encrypt = async (
    data: string,
    key: CryptoKey
): Promise<string> => {
    const iv = crypto.getRandomValues(new Uint8Array(12));
    const encrypted = await crypto.subtle.encrypt(
        { name: "AES-GCM", iv },
        key,
        new TextEncoder().encode(data)
    );

    return JSON.stringify({
        iv: Buffer.from(iv).toString("hex"),
        data: Buffer.from(encrypted).toString("hex"),
    });
};

/**
 * @function decrypt
 * @description Decrypts data that was encrypted using the encrypt function
 *
 * This function uses the Web Crypto API to decrypt session data that was previously
 * encrypted with the corresponding encrypt function. It expects the input to be a
 * JSON string containing both the encrypted data and the initialization vector (IV)
 * that was used for encryption.
 *
 * @param {string} encryptedData - The encrypted data JSON string containing both the
 *                                encrypted data and the initialization vector (IV)
 *                                as hex strings
 * @param {CryptoKey} key - The AES-GCM decryption key to use (same key used for encryption)
 *
 * @returns {Promise<string>} A promise that resolves to the decrypted data as a string
 *
 * @throws Will throw an error if decryption fails, which may happen if the encryption
 *        key is incorrect or the data has been tampered with
 */
export const decrypt = async (
    encryptedData: string,
    key: CryptoKey
): Promise<string> => {
    const { iv, data } = JSON.parse(encryptedData);
    const decrypted = await crypto.subtle.decrypt(
        { name: "AES-GCM", iv: Buffer.from(iv, "hex") },
        key,
        Buffer.from(data, "hex")
    );

    return new TextDecoder().decode(decrypted);
};

import { abstractTestnet } from "viem/chains";
import type { AbstractClient } from "@abstract-foundation/agw-client";
import type { Address } from "viem";
import { chain } from "@/config/chain";
import { clearStoredSession } from "./clear-stored-session-key";

/**
 * @function validateSession
 * @description Checks if a session is valid by querying the session validator contract

This function verifies whether a session is still valid (active) by calling the
sessionStatus function on the Abstract Global Wallet session validator contract.
If the session is found to be invalid (expired, closed, or non-existent), it
automatically cleans up the invalid session data.

The validation is performed on-chain by checking the status of the session hash
for the given wallet address. The status is mapped to the SessionStatus enum,
where Active (1) indicates a valid session.
 * @param {Address} address - The wallet address that owns the session
 * @param {string} sessionHash - The hash of the session to validate
 * @returns {Promise<boolean>} A promise that resolves to a boolean indicating whether
the session is valid (true) or not (false)
 */
export const validateSession = async (
    abstractClient: AbstractClient,
    address: Address,
    sessionHash: `0x${string}`
): Promise<boolean> => {
    try {
        const status = await abstractClient.getSessionStatus(sessionHash);

        // On Abstract testnet, any session is allowed, so we skip the check
        // However, on mainnet, we need to check if the session is both whitelisted and active.
        const isValid =
            status === SessionStatus.Active ||
            (chain === abstractTestnet && status === SessionStatus.NotInitialized);

        if (!isValid) {
            clearStoredSession(address);
        }

        return isValid;
    } catch (error) {
        console.error("Failed to validate session:", error);
        return false;
    }
};

/**
 * @enum {number} SessionStatus
 * @description Represents the possible statuses of an Abstract Global Wallet session
 *
 * This enum maps to the SessionKeyPolicyRegistry.Status values.
 * It's used to determine if a session is valid and can be used to submit transactions on behalf of the wallet.
 */
enum SessionStatus {
    /**
     * Session has not been initialized or does not exist
     */
    NotInitialized = 0,

    /**
     * Session is active and can be used to submit transactions
     */
    Active = 1,

    /**
     * Session has been manually closed/revoked by the wallet owner
     */
    Closed = 2,

    /**
     * Session has expired (exceeded its expiresAt timestamp)
     */
    Expired = 3,
}

export default SessionStatus;