import { useCallback, useEffect, useState, useMemo } from '@wordpress/element';
import { dispatch, select as selectData, useSelect } from '@wordpress/data';
import { invert, isEqual } from 'lodash';
import {
	getCurrentOrder,
	getCurrentOrderItems,
	getCurrentOrderShipments,
	getFirstSelectableOriginAddress,
	getDefaultOutboundAddress,
	getDefaultReturnAddress,
	getSubItems,
	getShipmentsAutogeneratedFromLabels,
	getStoreOrigin,
	getShipmentDefaultDates,
	SHIPMENT_TYPE,
} from 'utils';
import {
	LabelShipmentIdMap,
	Label,
	OriginAddress,
	ReturnShipmentInfo,
	SelectedRates,
	ShipmentItem,
	ShipmentDate,
	ShipmentType,
} from 'types';
import { addressStore } from 'data/address';
import { labelPurchaseStore } from 'data/label-purchase';
import { LABEL_PURCHASE_STATUS } from 'data/constants';

export function useShipmentState() {
	const [ currentShipmentId, setCurrentShipmentId ] = useState( '0' );
	const [ shipments, updateShipments ] = useState<
		Record< string, ShipmentItem[] >
	>( getCurrentOrderShipments() );
	const [ returnShipments, setReturnShipments ] = useState<
		Record< string, ReturnShipmentInfo >
	>( {} );
	const [ selections, setSelection ] = useState<
		Record< string, ShipmentItem[] >
	>( {
		0: [],
	} );

	const [ shipmentOrigins, setShipmentOrigins ] = useState<
		Record< string, OriginAddress | undefined >
	>( {
		0: getFirstSelectableOriginAddress(),
	} );

	// The most recently purchased label, that has not been refunded.
	const activePurchasedLabel =
		selectData( labelPurchaseStore ).getPurchasedLabel( currentShipmentId );

	const [ shipmentDates, setShipmentDates ] = useState<
		Record< string, ShipmentDate< Date > >
	>( {
		0: getShipmentDefaultDates( '0', activePurchasedLabel ),
	} );

	const findOriginAddressById = ( originId: string ) => {
		try {
			const store = selectData( addressStore );
			if ( store && typeof store.getOriginAddresses === 'function' ) {
				const origins: OriginAddress[] = store.getOriginAddresses();
				return origins?.find( ( a ) => a.id === originId );
			}
		} catch {
			// Handle cases where store or method doesn't exist
		}
		return undefined;
	};

	const isReturnShipment = useCallback(
		( shipmentKey: string ): boolean => {
			return !! returnShipments[ shipmentKey ]?.isReturn;
		},
		[ returnShipments ]
	);

	const getDefaultAddressForShipmentType = useCallback(
		( shipmentKey: string ): OriginAddress | undefined => {
			if ( isReturnShipment( shipmentKey ) ) {
				return getDefaultReturnAddress();
			}
			return getDefaultOutboundAddress();
		},
		[ isReturnShipment ]
	);

	const setShipmentOrigin = useCallback(
		( originId: string ) => {
			const origin = findOriginAddressById( originId );

			if (
				! origin ||
				( origin &&
					isEqual( shipmentOrigins[ currentShipmentId ], origin ) )
			) {
				return;
			}

			setShipmentOrigins( ( prevState ) => ( {
				...prevState,
				[ currentShipmentId ]: origin,
			} ) );
		},
		[ currentShipmentId, shipmentOrigins ]
	);

	const purchasedLabelOrigin = useSelect(
		( select ) =>
			select( labelPurchaseStore ).getLabelOrigins( currentShipmentId ),
		[ currentShipmentId ]
	);

	const purchasedLabelDestination = useSelect(
		( select ) =>
			select( labelPurchaseStore ).getLabelDestinations(
				currentShipmentId
			),
		// eslint-disable-next-line react-hooks/exhaustive-deps -- we want this to update when the shipmentId or activePurchasedLabel changes
		[ currentShipmentId, activePurchasedLabel ]
	);

	const orderDestination = useSelect(
		( select ) => select( addressStore ).getOrderDestination(),
		// order destination is not dependent on the shipmentId or activePurchasedLabel changes
		[]
	);

	useEffect( () => {
		// Fetching the origin and destination addresses for the most recently purchased label doesn't check
		// if it has been refunded or not, so we check for "activePurchasedLabel" as well.
		// older implementations of purchasedLabelOrigin don't have the id property, so we check for it.
		if (
			activePurchasedLabel &&
			[
				LABEL_PURCHASE_STATUS.PURCHASED,
				LABEL_PURCHASE_STATUS.PURCHASE_IN_PROGRESS,
			].includes( activePurchasedLabel.status ) &&
			purchasedLabelOrigin?.id
		) {
			setShipmentOrigin( purchasedLabelOrigin.id );
		} else if ( ! shipmentOrigins[ currentShipmentId ] ) {
			const defaultAddress =
				getDefaultAddressForShipmentType( currentShipmentId );
			if ( defaultAddress ) {
				setShipmentOrigin( defaultAddress.id );
			}
		}
	}, [
		currentShipmentId,
		activePurchasedLabel,
		purchasedLabelOrigin,
		purchasedLabelDestination,
		shipments,
		setShipmentOrigin,
		orderDestination,
		shipmentOrigins,
		getDefaultAddressForShipmentType,
	] );

	const [ labelShipmentIdsToUpdate, setLabelShipmentIdsToUpdate ] =
		useState< LabelShipmentIdMap >( {} );

	const getShipmentWeight = useCallback(
		() =>
			shipments[ currentShipmentId ]?.reduce(
				( acc, { weight, quantity } ) =>
					acc + Number( weight || 0 ) * Number( quantity ),
				0
			),
		[ shipments, currentShipmentId ]
	);

	const resetSelections = ( shipmentIds: string[] ) => {
		setSelection(
			shipmentIds.reduce(
				( acc, key ) => ( { ...acc, [ key ]: [] } ),
				{}
			)
		);
	};

	const getShipmentItems = useCallback(
		( shipmentId = currentShipmentId ) => shipments[ shipmentId ],
		[ shipments, currentShipmentId ]
	);

	const getSelectionItems = useCallback(
		( shipmentId = currentShipmentId ) => selections[ shipmentId ],
		[ selections, currentShipmentId ]
	);

	const getShipmentOrigin = useCallback( (): OriginAddress => {
		if (
			activePurchasedLabel &&
			purchasedLabelOrigin &&
			[
				LABEL_PURCHASE_STATUS.PURCHASED,
				LABEL_PURCHASE_STATUS.PURCHASE_IN_PROGRESS,
			].includes( activePurchasedLabel.status )
		) {
			const foundOrigin = findOriginAddressById(
				purchasedLabelOrigin.id
			);
			if ( foundOrigin ) {
				return foundOrigin;
			}
		}

		const selectedOrigin =
			shipmentOrigins[ currentShipmentId ] ??
			getDefaultAddressForShipmentType( currentShipmentId );
		// If no origin address is available, return a default empty address
		if ( ! selectedOrigin ) {
			const storeOrigin = getStoreOrigin();
			return {
				id: '',
				name: '',
				address: '',
				address1: '',
				address2: '',
				city: '',
				state: storeOrigin.state,
				postcode: '',
				country: storeOrigin.country,
				phone: '',
				email: '',
				firstName: '',
				lastName: '',
				company: '',
				isVerified: false,
				defaultAddress: false,
			};
		}

		return selectedOrigin;
	}, [
		activePurchasedLabel,
		currentShipmentId,
		purchasedLabelOrigin,
		shipmentOrigins,
		getDefaultAddressForShipmentType,
	] );

	/**
	 * Returns the ship from address recorded at the time of purchase.
	 * This can differ from the current ship from address if the user has changed it.
	 * If the user has changed the address, the id will remain the same, but the address will be different.
	 */
	const getShipmentPurchaseOrigin = () =>
		activePurchasedLabel &&
		purchasedLabelOrigin &&
		activePurchasedLabel.status === LABEL_PURCHASE_STATUS.PURCHASED
			? purchasedLabelOrigin
			: null;

	const getShipmentDestination = useCallback( () => {
		if (
			activePurchasedLabel &&
			activePurchasedLabel.status === LABEL_PURCHASE_STATUS.PURCHASED &&
			purchasedLabelDestination
		) {
			return purchasedLabelDestination;
		}

		return orderDestination;
	}, [ activePurchasedLabel, orderDestination, purchasedLabelDestination ] );

	const setShipments = (
		newShipments: Record< string, ShipmentItem[] >,
		updatedShipmentIds?: LabelShipmentIdMap
	) => {
		if ( updatedShipmentIds ) {
			setLabelShipmentIdsToUpdate( updatedShipmentIds );
			dispatch( labelPurchaseStore ).stageLabelsNewShipmentIds(
				updatedShipmentIds
			);
		}

		updateShipments( newShipments );
	};

	const revertLabelShipmentIdsToUpdate = () => {
		dispatch( labelPurchaseStore ).stageLabelsNewShipmentIds(
			invert( labelShipmentIdsToUpdate )
		);
		setLabelShipmentIdsToUpdate( {} );
	};

	const hasVariations = useMemo(
		() =>
			Object.values( shipments )
				.flat()
				.some( ( item ) => {
					return !! item.variation?.length;
				} ),
		[ shipments ]
	);

	const hasMultipleShipments = useMemo(
		() => Object.values( shipments ).length > 1,
		[ shipments ]
	);

	const isExtraLabelPurchaseValid = () => {
		return selections[ currentShipmentId ]?.length > 0;
	};

	const resetShipmentAndSelection = useCallback( () => {
		const updatedShipments = {
			...shipments,
			[ currentShipmentId ]: undefined,
		};
		delete updatedShipments[ currentShipmentId ];

		const order = getCurrentOrder();

		dispatch( labelPurchaseStore ).updateShipments( {
			shipments: updatedShipments,
			orderId: `${ order?.id }`,
			shipmentIdsToUpdate: {},
		} );
		const orderItems = getCurrentOrderItems();
		const shipmentItems = orderItems.map( ( orderItem ) => ( {
			...orderItem,
			subItems: getSubItems( orderItem as ShipmentItem ),
		} ) );
		updateShipments( ( currentShipments ) => ( {
			...currentShipments,
			[ currentShipmentId ]: shipmentItems,
		} ) );
		setSelection( {
			...selections,
			[ currentShipmentId ]: shipmentItems,
		} );
	}, [ currentShipmentId, selections, shipments ] );

	const isShipmentAutogeneratedFromLabel = useCallback(
		() =>
			getShipmentsAutogeneratedFromLabels().includes(
				Number( currentShipmentId )
			),
		[ currentShipmentId ]
	);

	const setCurrentShipmentDate = useCallback(
		( shippingDate: Date, estimatedDeliveryDate?: Date ) => {
			setShipmentDates( ( prevState ) => ( {
				...prevState,
				[ currentShipmentId ]: { shippingDate, estimatedDeliveryDate },
			} ) );
		},
		[ currentShipmentId ]
	);

	const getCurrentShipmentDate = useCallback( () => {
		return shipmentDates[ currentShipmentId ];
	}, [ currentShipmentId, shipmentDates ] );

	/**
	 * Load return shipment data from purchased labels on initialization.
	 * This ensures parentShipmentId persists across page reloads.
	 */
	useEffect( () => {
		try {
			const store = selectData( labelPurchaseStore );
			if ( store && typeof store.getPurchasedLabels === 'function' ) {
				const purchasedLabels: Record< string, Label | undefined > =
					store.getPurchasedLabels();
				const returnShipmentsFromLabels: Record<
					string,
					ReturnShipmentInfo
				> = {};

				// Extract return shipment info from purchased labels
				Object.entries( purchasedLabels ).forEach(
					( [ shipmentId, label ] ) => {
						if ( label?.isReturn ) {
							returnShipmentsFromLabels[ shipmentId ] = {
								isReturn: true,
								parentShipmentId: label.parentShipmentId,
							};
						}
					}
				);

				// Also check selected rates for any in-progress return shipments
				if ( typeof store.getSelectedRates === 'function' ) {
					const selectedRates: SelectedRates | undefined =
						store.getSelectedRates();
					if ( selectedRates ) {
						Object.entries( selectedRates ).forEach(
							( [ shipmentKey, rateData ] ) => {
								if (
									rateData?.rate?.isReturn &&
									! returnShipmentsFromLabels[ shipmentKey ]
								) {
									// Only add from rates if not already in purchased labels
									returnShipmentsFromLabels[ shipmentKey ] = {
										isReturn: true,
									};
								}
							}
						);
					}
				}

				setReturnShipments( returnShipmentsFromLabels );
			}
		} catch {
			// Silently handle cases where the store or method doesn't exist (e.g., in tests)
		}
	}, [] );

	const getCurrentShipmentIsReturn = () => {
		return isReturnShipment( currentShipmentId );
	};

	const getCurrentShipmentParentId = () => {
		return returnShipments[ currentShipmentId ]?.parentShipmentId;
	};

	// Only set default dates if the current shipment doesn't have dates defined
	useEffect( () => {
		if (
			shipmentDates[ currentShipmentId ]?.shippingDate ||
			shipmentDates[ currentShipmentId ]?.estimatedDeliveryDate
		) {
			return;
		}
		setShipmentDates( ( prevState ) => ( {
			...prevState,
			[ currentShipmentId ]: getShipmentDefaultDates(
				currentShipmentId,
				activePurchasedLabel
			),
		} ) );
	}, [ currentShipmentId, shipmentDates, activePurchasedLabel ] );

	const getShipmentType = useCallback(
		( shipmentKey: string ): ShipmentType => {
			return isReturnShipment( shipmentKey )
				? SHIPMENT_TYPE.RETURN
				: SHIPMENT_TYPE.OUTBOUND;
		},
		[ isReturnShipment ]
	);

	return {
		shipments,
		setShipments,
		getShipmentWeight,
		resetSelections,
		selections,
		setSelection,
		currentShipmentId,
		setCurrentShipmentId,
		getShipmentItems,
		getSelectionItems,
		getShipmentOrigin,
		setShipmentOrigin,
		getShipmentDestination,
		revertLabelShipmentIdsToUpdate,
		labelShipmentIdsToUpdate,
		getShipmentPurchaseOrigin,
		hasVariations,
		hasMultipleShipments,
		isExtraLabelPurchaseValid,
		resetShipmentAndSelection,
		isShipmentAutogeneratedFromLabel,
		setCurrentShipmentDate,
		getCurrentShipmentDate,
		getCurrentShipmentIsReturn,
		getCurrentShipmentParentId,
		setReturnShipments,
		returnShipments,
		isReturnShipment,
		getShipmentType,
	};
}