import * as Baseten from "../providers/baseten.js";
import * as Cerebras from "../providers/cerebras.js";
import * as Cohere from "../providers/cohere.js";
import * as DeepInfra from "../providers/deepinfra.js";
import * as FalAI from "../providers/fal-ai.js";
import * as FeatherlessAI from "../providers/featherless-ai.js";
import * as Fireworks from "../providers/fireworks-ai.js";
import * as Groq from "../providers/groq.js";
import * as HFInference from "../providers/hf-inference.js";
import * as Novita from "../providers/novita.js";
import * as Nscale from "../providers/nscale.js";
import * as OpenAI from "../providers/openai.js";
import * as OvhCloud from "../providers/ovhcloud.js";
import * as PublicAI from "../providers/publicai.js";
import type {
	AudioClassificationTaskHelper,
	AudioToAudioTaskHelper,
	AutomaticSpeechRecognitionTaskHelper,
	ConversationalTaskHelper,
	DocumentQuestionAnsweringTaskHelper,
	FeatureExtractionTaskHelper,
	FillMaskTaskHelper,
	ImageClassificationTaskHelper,
	ImageSegmentationTaskHelper,
	ImageToImageTaskHelper,
	ImageToTextTaskHelper,
	ImageToVideoTaskHelper,
	ImageTextToImageTaskHelper,
	ImageTextToVideoTaskHelper,
	ObjectDetectionTaskHelper,
	QuestionAnsweringTaskHelper,
	SentenceSimilarityTaskHelper,
	SummarizationTaskHelper,
	TableQuestionAnsweringTaskHelper,
	TabularClassificationTaskHelper,
	TabularRegressionTaskHelper,
	TaskProviderHelper,
	TextClassificationTaskHelper,
	TextGenerationTaskHelper,
	TextToAudioTaskHelper,
	TextToImageTaskHelper,
	TextToSpeechTaskHelper,
	TextToVideoTaskHelper,
	TokenClassificationTaskHelper,
	TranslationTaskHelper,
	VisualQuestionAnsweringTaskHelper,
	ZeroShotClassificationTaskHelper,
	ZeroShotImageClassificationTaskHelper,
} from "../providers/providerHelper.js";
import * as Replicate from "../providers/replicate.js";
import * as Scaleway from "../providers/scaleway.js";
import * as Together from "../providers/together.js";
import * as Wavespeed from "../providers/wavespeed.js";
import * as Zai from "../providers/zai-org.js";
import type { InferenceProvider, InferenceProviderOrPolicy, InferenceTask } from "../types.js";
import { InferenceClientInputError } from "../errors.js";

export const PROVIDERS: Record<InferenceProvider, Partial<Record<InferenceTask, TaskProviderHelper>>> = {
	baseten: {
		conversational: new Baseten.BasetenConversationalTask(),
	},
	cerebras: {
		conversational: new Cerebras.CerebrasConversationalTask(),
	},
	cohere: {
		conversational: new Cohere.CohereConversationalTask(),
	},
	deepinfra: {
		"automatic-speech-recognition": new DeepInfra.DeepInfraAutomaticSpeechRecognitionTask(),
		conversational: new DeepInfra.DeepInfraConversationalTask(),
		"text-generation": new DeepInfra.DeepInfraTextGenerationTask(),
	},
	"fal-ai": {
		"audio-to-audio": new FalAI.FalAIAudioToAudioTask(),
		"automatic-speech-recognition": new FalAI.FalAIAutomaticSpeechRecognitionTask(),
		"image-text-to-image": new FalAI.FalAIImageTextToImageTask(),
		"image-text-to-video": new FalAI.FalAIImageTextToVideoTask(),
		"image-to-image": new FalAI.FalAIImageToImageTask(),
		"image-segmentation": new FalAI.FalAIImageSegmentationTask(),
		"image-to-video": new FalAI.FalAIImageToVideoTask(),
		"text-to-audio": new FalAI.FalAITextToAudioTask(),
		"text-to-image": new FalAI.FalAITextToImageTask(),
		"text-to-speech": new FalAI.FalAITextToSpeechTask(),
		"text-to-video": new FalAI.FalAITextToVideoTask(),
	},
	"featherless-ai": {
		conversational: new FeatherlessAI.FeatherlessAIConversationalTask(),
		"text-generation": new FeatherlessAI.FeatherlessAITextGenerationTask(),
	},
	"hf-inference": {
		"text-to-image": new HFInference.HFInferenceTextToImageTask(),
		conversational: new HFInference.HFInferenceConversationalTask(),
		"text-generation": new HFInference.HFInferenceTextGenerationTask(),
		"text-classification": new HFInference.HFInferenceTextClassificationTask(),
		"question-answering": new HFInference.HFInferenceQuestionAnsweringTask(),
		"audio-classification": new HFInference.HFInferenceAudioClassificationTask(),
		"automatic-speech-recognition": new HFInference.HFInferenceAutomaticSpeechRecognitionTask(),
		"fill-mask": new HFInference.HFInferenceFillMaskTask(),
		"feature-extraction": new HFInference.HFInferenceFeatureExtractionTask(),
		"image-classification": new HFInference.HFInferenceImageClassificationTask(),
		"image-segmentation": new HFInference.HFInferenceImageSegmentationTask(),
		"document-question-answering": new HFInference.HFInferenceDocumentQuestionAnsweringTask(),
		"image-to-text": new HFInference.HFInferenceImageToTextTask(),
		"object-detection": new HFInference.HFInferenceObjectDetectionTask(),
		"audio-to-audio": new HFInference.HFInferenceAudioToAudioTask(),
		"zero-shot-image-classification": new HFInference.HFInferenceZeroShotImageClassificationTask(),
		"zero-shot-classification": new HFInference.HFInferenceZeroShotClassificationTask(),
		"image-to-image": new HFInference.HFInferenceImageToImageTask(),
		"sentence-similarity": new HFInference.HFInferenceSentenceSimilarityTask(),
		"table-question-answering": new HFInference.HFInferenceTableQuestionAnsweringTask(),
		"tabular-classification": new HFInference.HFInferenceTabularClassificationTask(),
		"text-to-speech": new HFInference.HFInferenceTextToSpeechTask(),
		"token-classification": new HFInference.HFInferenceTokenClassificationTask(),
		translation: new HFInference.HFInferenceTranslationTask(),
		summarization: new HFInference.HFInferenceSummarizationTask(),
		"visual-question-answering": new HFInference.HFInferenceVisualQuestionAnsweringTask(),
		"tabular-regression": new HFInference.HFInferenceTabularRegressionTask(),
		"text-to-audio": new HFInference.HFInferenceTextToAudioTask(),
	},
	"fireworks-ai": {
		conversational: new Fireworks.FireworksConversationalTask(),
	},
	groq: {
		conversational: new Groq.GroqConversationalTask(),
		"text-generation": new Groq.GroqTextGenerationTask(),
	},
	novita: {
		conversational: new Novita.NovitaConversationalTask(),
		"text-generation": new Novita.NovitaTextGenerationTask(),
		"text-to-video": new Novita.NovitaTextToVideoTask(),
	},
	nscale: {
		"text-to-image": new Nscale.NscaleTextToImageTask(),
		conversational: new Nscale.NscaleConversationalTask(),
	},
	openai: {
		conversational: new OpenAI.OpenAIConversationalTask(),
	},
	ovhcloud: {
		conversational: new OvhCloud.OvhCloudConversationalTask(),
		"text-generation": new OvhCloud.OvhCloudTextGenerationTask(),
	},
	publicai: {
		conversational: new PublicAI.PublicAIConversationalTask(),
	},
	replicate: {
		"text-to-image": new Replicate.ReplicateTextToImageTask(),
		"text-to-speech": new Replicate.ReplicateTextToSpeechTask(),
		"text-to-video": new Replicate.ReplicateTextToVideoTask(),
		"image-to-image": new Replicate.ReplicateImageToImageTask(),
		"automatic-speech-recognition": new Replicate.ReplicateAutomaticSpeechRecognitionTask(),
	},
	scaleway: {
		conversational: new Scaleway.ScalewayConversationalTask(),
		"text-generation": new Scaleway.ScalewayTextGenerationTask(),
		"feature-extraction": new Scaleway.ScalewayFeatureExtractionTask(),
	},
	together: {
		"text-to-image": new Together.TogetherTextToImageTask(),
		"image-to-image": new Together.TogetherImageToImageTask(),
		"text-to-video": new Together.TogetherTextToVideoTask(),
		"image-to-video": new Together.TogetherImageToVideoTask(),
		conversational: new Together.TogetherConversationalTask(),
		"text-generation": new Together.TogetherTextGenerationTask(),
		"feature-extraction": new Together.TogetherFeatureExtractionTask(),
		"text-to-speech": new Together.TogetherTextToSpeechTask(),
		"automatic-speech-recognition": new Together.TogetherAutomaticSpeechRecognitionTask(),
	},
	wavespeed: {
		"text-to-image": new Wavespeed.WavespeedAITextToImageTask(),
		"text-to-video": new Wavespeed.WavespeedAITextToVideoTask(),
		"image-to-image": new Wavespeed.WavespeedAIImageToImageTask(),
		"image-to-video": new Wavespeed.WavespeedAIImageToVideoTask(),
		"image-text-to-image": new Wavespeed.WavespeedAIImageTextToImageTask(),
		"image-text-to-video": new Wavespeed.WavespeedAIImageTextToVideoTask(),
	},
	"zai-org": {
		conversational: new Zai.ZaiConversationalTask(),
		"text-to-image": new Zai.ZaiTextToImageTask(),
		"image-to-text": new Zai.ZaiImageToTextTask(),
	},
};

/**
 * Get provider helper instance by name and task
 */
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "text-to-image",
): TextToImageTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "conversational",
): ConversationalTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "text-generation",
): TextGenerationTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "text-to-speech",
): TextToSpeechTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "text-to-audio",
): TextToAudioTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "automatic-speech-recognition",
): AutomaticSpeechRecognitionTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "text-to-video",
): TextToVideoTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "text-classification",
): TextClassificationTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "question-answering",
): QuestionAnsweringTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "audio-classification",
): AudioClassificationTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "audio-to-audio",
): AudioToAudioTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "fill-mask",
): FillMaskTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "feature-extraction",
): FeatureExtractionTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "image-classification",
): ImageClassificationTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "image-segmentation",
): ImageSegmentationTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "document-question-answering",
): DocumentQuestionAnsweringTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "image-to-text",
): ImageToTextTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "object-detection",
): ObjectDetectionTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "zero-shot-image-classification",
): ZeroShotImageClassificationTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "zero-shot-classification",
): ZeroShotClassificationTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "image-to-image",
): ImageToImageTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "image-to-video",
): ImageToVideoTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "image-text-to-image",
): ImageTextToImageTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "image-text-to-video",
): ImageTextToVideoTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "sentence-similarity",
): SentenceSimilarityTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "table-question-answering",
): TableQuestionAnsweringTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "tabular-classification",
): TabularClassificationTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "tabular-regression",
): TabularRegressionTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "token-classification",
): TokenClassificationTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "translation",
): TranslationTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "summarization",
): SummarizationTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: "visual-question-answering",
): VisualQuestionAnsweringTaskHelper & TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: InferenceTask | undefined,
): TaskProviderHelper;
export function getProviderHelper(
	provider: InferenceProviderOrPolicy,
	task: InferenceTask | undefined,
): TaskProviderHelper {
	if ((provider === "hf-inference" && !task) || provider === "auto") {
		return new HFInference.HFInferenceTask();
	}
	if (!task) {
		throw new InferenceClientInputError(
			"you need to provide a task name when using an external provider, e.g. 'text-to-image'",
		);
	}
	if (!(provider in PROVIDERS)) {
		throw new InferenceClientInputError(
			`Provider '${provider}' not supported. Available providers: ${Object.keys(PROVIDERS)}`,
		);
	}
	const providerTasks = PROVIDERS[provider];
	if (!providerTasks || !(task in providerTasks)) {
		throw new InferenceClientInputError(
			`Task '${task}' not supported for provider '${provider}'. Available tasks: ${Object.keys(providerTasks ?? {})}`,
		);
	}
	return providerTasks[task] as TaskProviderHelper;
}