"use client";

import React, { createContext, useContext, useState } from "react";

interface TrialVariantSelectionContextValue {
  selectedVariantId: string | null;
  setSelectedVariantId: (variantId: string | null) => void;
}

const TrialVariantSelectionContext =
  createContext<TrialVariantSelectionContextValue | null>(null);

const STORAGE_KEY = 'trial_variant_selection';

export function TrialVariantSelectionProvider({
  children,
}: {
  children: React.ReactNode;
}) {
  // Initialize from sessionStorage if available
  const [selectedVariantId, setSelectedVariantId] = useState<string | null>(() => {
    if (typeof window === 'undefined') return null;
    const stored = sessionStorage.getItem(STORAGE_KEY);
    return stored || null;
  });

  const wrappedSetSelectedVariantId = (id: string | null) => {
    setSelectedVariantId(id);
    
    // Persist to sessionStorage
    if (typeof window !== 'undefined') {
      if (id) {
        sessionStorage.setItem(STORAGE_KEY, id);
      } else {
        sessionStorage.removeItem(STORAGE_KEY);
      }
    }
  };
  
  return (
    <TrialVariantSelectionContext.Provider
      value={{ selectedVariantId, setSelectedVariantId: wrappedSetSelectedVariantId }}
    >
      {children}
    </TrialVariantSelectionContext.Provider>
  );
}

export function useTrialVariantSelection() {
  const context = useContext(TrialVariantSelectionContext);
  if (!context) {
    throw new Error(
      "useTrialVariantSelection must be used within TrialVariantSelectionProvider"
    );
  }
  return context;
}