frontend/src/lib/ImageContext.tsx

import { createContext, useContext, useState, ReactNode, useCallback } from 'react';
import { ClassificationResult } from './api';

export interface ImageState {
    // Current image
    file: File | null;
    preview: string | null;
    processedImage: string | null;

    // Classification results
    classificationResults: ClassificationResult[] | null;

    // View (corrected takes priority)
    predictedView: string | null;
    correctedView: string | null;

    // Session ID for feedback
    sessionId: string;
}

interface ImageContextType extends ImageState {
    // Actions
    setFile: (file: File | null, preview: string | null) => void;
    setClassificationResults: (results: ClassificationResult[] | null, processedImage?: string | null) => void;
    setCorrectedView: (view: string | null) => void;
    resetState: () => void;

    // Computed
    currentView: string | null; // Returns correctedView if set, otherwise topPrediction
    isViewGAEligible: boolean;
}

// GA-eligible views (match keys from prompt_fetal_view.json)
export const GA_ELIGIBLE_VIEWS = [
    "brain",    // Head Circumference
    "abdomen",  // Abdominal Circumference
    "femur",    // Femur Length
];

// Display names for GA biometry types
export const GA_BIOMETRY_LABELS: Record<string, string> = {
    "brain": "Head Circumference (HC)",
    "abdomen": "Abdominal Circumference (AC)",
    "femur": "Femur Length (FL)",
};

const generateSessionId = () => {
    return Math.random().toString(36).substring(2, 10);
};

const initialState: ImageState = {
    file: null,
    preview: null,
    processedImage: null,
    classificationResults: null,
    predictedView: null,
    correctedView: null,
    sessionId: generateSessionId(),
};

const ImageContext = createContext<ImageContextType | null>(null);

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

    const setFile = useCallback((file: File | null, preview: string | null) => {
        setState(prev => ({
            ...prev,
            file,
            preview,
            processedImage: null,
            classificationResults: null,
            predictedView: null,
            correctedView: null,
        }));
    }, []);

    const setClassificationResults = useCallback((
        results: ClassificationResult[] | null,
        processedImage: string | null = null
    ) => {
        setState(prev => ({
            ...prev,
            classificationResults: results,
            processedImage: processedImage || prev.processedImage,
            predictedView: results && results.length > 0 ? results[0].label : null,
        }));
    }, []);

    const setCorrectedView = useCallback((view: string | null) => {
        setState(prev => ({
            ...prev,
            correctedView: view,
        }));
    }, []);

    const resetState = useCallback(() => {
        setState({
            ...initialState,
            sessionId: generateSessionId(),
        });
    }, []);

    // Computed: current view (corrected takes priority)
    const currentView = state.correctedView || state.predictedView;

    // Computed: is view eligible for GA estimation
    const isViewGAEligible = currentView ? GA_ELIGIBLE_VIEWS.includes(currentView) : false;

    return (
        <ImageContext.Provider
            value={{
                ...state,
                setFile,
                setClassificationResults,
                setCorrectedView,
                resetState,
                currentView,
                isViewGAEligible,
            }}
        >
            {children}
        </ImageContext.Provider>
    );
}

export function useImageContext() {
    const context = useContext(ImageContext);
    if (!context) {
        throw new Error('useImageContext must be used within an ImageProvider');
    }
    return context;
}