import type { ComponentType, SVGProps } from "react";
import OpenAI from "@lobehub/icons/es/OpenAI";
import Anthropic from "@lobehub/icons/es/Anthropic";
import Google from "@lobehub/icons/es/Google";
import Gemini from "@lobehub/icons/es/Gemini";
import DeepSeek from "@lobehub/icons/es/DeepSeek";
import Meta from "@lobehub/icons/es/Meta";
import Mistral from "@lobehub/icons/es/Mistral";
import Claude from "@lobehub/icons/es/Claude";
import Qwen from "@lobehub/icons/es/Qwen";
import Zhipu from "@lobehub/icons/es/Zhipu";
import ChatGLM from "@lobehub/icons/es/ChatGLM";
import Kimi from "@lobehub/icons/es/Kimi";
import Minimax from "@lobehub/icons/es/Minimax";
import Moonshot from "@lobehub/icons/es/Moonshot";
import Stepfun from "@lobehub/icons/es/Stepfun";
import Yi from "@lobehub/icons/es/Yi";
import Doubao from "@lobehub/icons/es/Doubao";
import ByteDance from "@lobehub/icons/es/ByteDance";
import Baidu from "@lobehub/icons/es/Baidu";
import Hunyuan from "@lobehub/icons/es/Hunyuan";
import Tencent from "@lobehub/icons/es/Tencent";
import Spark from "@lobehub/icons/es/Spark";
import Baichuan from "@lobehub/icons/es/Baichuan";

type LobeIcon = ComponentType<SVGProps<SVGSVGElement>>;

/** Prefix rules for matching model IDs to icons. First match wins. */
const MODEL_PREFIXES: Array<[string, LobeIcon]> = [
  // OpenAI
  ["gpt-", OpenAI],
  ["o1-", OpenAI],
  ["o3-", OpenAI],
  ["o4-", OpenAI],
  ["chatgpt-", OpenAI],
  // Anthropic / Claude
  ["claude-", Claude],
  ["anthropic-", Anthropic],
  // Google
  ["gemini-", Gemini],
  ["gemma-", Google],
  ["palm-", Google],
  // DeepSeek
  ["deepseek-", DeepSeek],
  // Meta
  ["llama-", Meta],
  ["meta-", Meta],
  // Mistral
  ["mistral-", Mistral],
  ["codestral-", Mistral],
  ["mixtral-", Mistral],
  // Alibaba
  ["qwen-", Qwen],
  ["tongyi-", Qwen],
  // Zhipu / GLM
  ["glm-", Zhipu],
  ["zhipu-", Zhipu],
  ["chatglm-", ChatGLM],
  // Moonshot / Kimi
  ["kimi-", Kimi],
  ["moonshot-", Moonshot],
  // MiniMax
  ["minimax-", Minimax],
  ["abab-", Minimax],
  // Stepfun
  ["step-", Stepfun],
  // 01.AI / Yi
  ["yi-", Yi],
  // ByteDance / Doubao
  ["doubao-", Doubao],
  ["skylark-", ByteDance],
  // Baidu
  ["ernie-", Baidu],
  ["baidu-", Baidu],
  // Tencent
  ["hunyuan-", Hunyuan],
  ["tencent-", Tencent],
  // iFlytek
  ["spark-", Spark],
  ["iflytek-", Spark],
  // Baichuan
  ["baichuan-", Baichuan],
];

/** Provider name → icon (case-insensitive key). */
const PROVIDER_ICONS: Record<string, LobeIcon> = {
  openai: OpenAI,
  anthropic: Anthropic,
  google: Google,
  gemini: Gemini,
  deepseek: DeepSeek,
  meta: Meta,
  mistral: Mistral,
  claude: Claude,
  qwen: Qwen,
  zhipu: Zhipu,
  glm: Zhipu,
  chatglm: ChatGLM,
  kimi: Kimi,
  minimax: Minimax,
  moonshot: Moonshot,
  stepfun: Stepfun,
  yi: Yi,
  doubao: Doubao,
  bytedance: ByteDance,
  baidu: Baidu,
  hunyuan: Hunyuan,
  tencent: Tencent,
  spark: Spark,
  iflytek: Spark,
  baichuan: Baichuan,
};

export function getModelIcon(
  modelName: string,
  providerName?: string,
): LobeIcon | null {
  const lower = modelName.toLowerCase().trim();
  for (const [prefix, icon] of MODEL_PREFIXES) {
    if (lower.startsWith(prefix)) return icon;
  }
  if (providerName) return getProviderIcon(providerName);
  return null;
}

export function getProviderIcon(providerName: string): LobeIcon | null {
  const key = providerName.toLowerCase().trim();
  if (PROVIDER_ICONS[key]) return PROVIDER_ICONS[key];
  for (const [k, icon] of Object.entries(PROVIDER_ICONS)) {
    if (key.includes(k)) return icon;
  }
  return null;
}
