import { IPromptAgent } from '../types';
import { PhotorealisticAgent } from './PhotorealisticAgent';
import { IllustrationAgent } from './IllustrationAgent';
import { MinimalistAgent } from './MinimalistAgent';
import { StickerAgent } from './StickerAgent';
import { ProductAgent } from './ProductAgent';
import { ComicAgent } from './ComicAgent';
import { GeneralAgent } from './GeneralAgent';

type AgentConstructor = new (apiKey: string) => IPromptAgent;

const AGENT_REGISTRY: Record<string, AgentConstructor> = {
  photorealistic: PhotorealisticAgent,
  illustration: IllustrationAgent,
  minimalist: MinimalistAgent,
  sticker: StickerAgent,
  product: ProductAgent,
  comic: ComicAgent,
  general: GeneralAgent,
};

export function getAgent(apiKey: string, template?: string): IPromptAgent {
  if (!template) {
    return new GeneralAgent(apiKey);
  }

  const AgentClass = AGENT_REGISTRY[template];
  if (!AgentClass) {
    console.warn(`Unknown template "${template}", falling back to GeneralAgent`);
    return new GeneralAgent(apiKey);
  }

  return new AgentClass(apiKey);
}

export {
  PhotorealisticAgent,
  IllustrationAgent,
  MinimalistAgent,
  StickerAgent,
  ProductAgent,
  ComicAgent,
  GeneralAgent,
};
