import { NodeType } from '../../types/node';
import { NodeExecutionStrategy } from './types';
import { ImageStrategy } from './strategies/ImageStrategy';
import { VideoStrategy } from './strategies/VideoStrategy';
import { AudioStrategy } from './strategies/AudioStrategy';
import { VideoAnalyzerStrategy } from './strategies/AnalysisStrategy';
import { ImageEditorStrategy, PromptGeneratorStrategy, LyricsStrategy, GenericSuccessStrategy } from './strategies/DefaultStrategy';

const STRATEGY_MAP: Record<string, NodeExecutionStrategy> = {
    [NodeType.IMAGE_GENERATOR]: new ImageStrategy(),
    [NodeType.VIDEO_GENERATOR]: new VideoStrategy(),
    [NodeType.AUDIO_GENERATOR]: new AudioStrategy(),
    [NodeType.LYRICS_GENERATOR]: new LyricsStrategy(),
    [NodeType.VIDEO_ANALYZER]: new VideoAnalyzerStrategy(),
    [NodeType.IMAGE_EDITOR]: new ImageEditorStrategy(),
    [NodeType.PROMPT_GENERATOR]: new PromptGeneratorStrategy(),
};

const DEFAULT_STRATEGY = new GenericSuccessStrategy();

export function getExecutionStrategy(type: string): NodeExecutionStrategy {
    return STRATEGY_MAP[type] || DEFAULT_STRATEGY;
}

export * from './types';
