import React, { createContext, useContext, useState, ReactNode } from 'react';
import { AstroData, Profile } from '../types';

export type ScreenId = 1 | 2 | 3 | 4 | 5;

export interface SessionState {
  currentScreen: ScreenId;
  step1_2_Data: (AstroData & { summary?: string }) | null;
  selectedLifeArea: string | null;
  selectedVargas: string[];
  step3_Data: Record<string, string>;
  step4_Data: {
    vimshottariDasha: string;
    rasiDasha: string;
    transits: string;
    ashtakavarga: string;
  } | null;
  modelPreferences: Record<string, string>;
  lastError: Record<string, string | null>;
  // Multi-user Profile state
  profiles: Profile[];
  activeProfileId: string | null;
}

interface AstroContextType {
  state: SessionState;
  setScreen: (screen: ScreenId) => void;
  setStep1Data: (data: AstroData) => void;
  setStep1Summary: (summary: string) => void;
  setSelectedLifeArea: (areaId: string | null) => void;
  setSelectedVargas: (vargas: string[]) => void;
  setStep3Data: (vargaKey: string, data: string) => void;
  setStep4Data: (data: SessionState['step4_Data']) => void;
  setModelPreference: (stepId: string, model: string) => void;
  setStepError: (stepId: string, error: string | null) => void;
  setProfiles: (profiles: Profile[]) => void;
  setActiveProfileId: (id: string | null) => void;
  resetSession: () => void;
}

const initialState: SessionState = {
  currentScreen: 1,
  step1_2_Data: null,
  selectedLifeArea: null,
  selectedVargas: ['D1', 'D9', 'D10'], // Default priority
  step3_Data: {},
  step4_Data: null,
  modelPreferences: {
    'screen2': 'gemini-2.5-flash',
    'screen3': 'gemini-2.5-flash',
    'screen4': 'gemini-2.5-flash',
  },
  lastError: {},
  profiles: [],
  activeProfileId: null,
};

const AstroContext = createContext<AstroContextType | undefined>(undefined);

export function AstroProvider({ children }: { children: ReactNode }) {
  const [state, setState] = useState<SessionState>(initialState);

  const setScreen = (screen: ScreenId) => {
    setState(prev => ({ ...prev, currentScreen: screen }));
    window.scrollTo({ top: 0, behavior: 'smooth' });
  };

  const setStep1Data = (data: AstroData) => {
    setState(prev => ({ ...prev, step1_2_Data: data }));
  };

  const setStep1Summary = (summary: string) => {
    setState(prev => ({
      ...prev,
      step1_2_Data: prev.step1_2_Data ? { ...prev.step1_2_Data, summary } : null
    }));
  };

  const setSelectedLifeArea = (areaId: string | null) => {
    setState(prev => ({ ...prev, selectedLifeArea: areaId }));
  };

  const setSelectedVargas = (vargas: string[]) => {
    setState(prev => ({ ...prev, selectedVargas: vargas }));
  };

  const setStep3Data = (vargaKey: string, data: string) => {
    setState(prev => ({
      ...prev,
      step3_Data: { ...prev.step3_Data, [vargaKey]: data }
    }));
  };

  const setStep4Data = (data: SessionState['step4_Data']) => {
    setState(prev => ({ ...prev, step4_Data: data }));
  };

  const setModelPreference = (stepId: string, model: string) => {
    setState(prev => ({
      ...prev,
      modelPreferences: { ...prev.modelPreferences, [stepId]: model }
    }));
  };

  const setStepError = (stepId: string, error: string | null) => {
    setState(prev => ({
      ...prev,
      lastError: { ...prev.lastError, [stepId]: error }
    }));
  };

  const setProfiles = (profiles: Profile[]) => {
    setState(prev => ({ ...prev, profiles }));
  };

  const setActiveProfileId = (id: string | null) => {
    setState(prev => ({ ...prev, activeProfileId: id }));
  };

  const resetSession = () => {
    if (window.confirm("Are you sure you want to reset all data and start a new session?")) {
      setState(initialState);
    }
  };

  return (
    <AstroContext.Provider 
      value={{ 
        state, 
        setScreen, 
        setStep1Data, 
        setStep1Summary,
        setSelectedLifeArea,
        setSelectedVargas, 
        setStep3Data, 
        setStep4Data, 
        setModelPreference,
        setStepError,
        setProfiles,
        setActiveProfileId,
        resetSession 
      }}
    >
      {children}
    </AstroContext.Provider>
  );
}


export function useAstroContext() {
  const context = useContext(AstroContext);
  if (context === undefined) {
    throw new Error('useAstroContext must be used within an AstroProvider');
  }
  return context;
}