import { create } from 'zustand';
import { Doubao, Qwen, ModelScope, OpenAI, Midjourney, Kling, Minimax, Flux, Runway, Vidu, Jimeng, Moonshot, AlibabaCloud, Anthropic, Gemini} from '@lobehub/icons';

// Icon mapping based on provider or explicit ID
const ICON_MAP: Record<string, any> = {
    'dashscope': Qwen,
    'volcengine': Doubao,
    'modelscope': ModelScope,
    'openai': OpenAI,
    'midjourney': Midjourney,
    'google': Gemini,
    'kling': Kling,
    'minimax': Minimax,
    'flux': Flux,
    'runway': Runway,
    'vidu': Vidu,
    'jimeng': Jimeng,
    'ali-videoenhan': AlibabaCloud,
    'moonshot':Moonshot,
    'anthropic': Anthropic,
};

export interface ModelCapability {
    // Generation modes
    text2video?: boolean;
    image2video?: boolean;
    video2video?: boolean;
    omni?: boolean;
    // Frame control
    supportsFirstFrame?: boolean;
    supportsLastFrame?: boolean;
    // Input types
    supportsMultiImage?: boolean;
    supportsMultiVideo?: boolean;
    supportsRefImage?: boolean;
    supportsReferenceUrls?: boolean;
    // Extensions
    supportsLora?: boolean;
    supportsAudio?: boolean;
    supportsSound?: boolean;
    supportsShotType?: boolean;
    supportsNegativePrompt?: boolean;
    supportsCameraControl?: boolean;
}

export interface ModelConfig {
    id: string;
    name: string;
    type: 'image' | 'video' | 'llm' | 'audio' | 'lyrics' | 'music';
    provider: string;
    is_default: boolean;
    enabled?: boolean;
    capabilities: ModelCapability;
    resolutions?: Record<string, any>;
    durations?: any;
    counts?: { min: number; max: number };
    voices?: Array<{ 
        id?: string; 
        name?: string; 
        value?: string; 
        label?: string;
        gender?: string;
        desc?: string;
    }>;
    icon?: any;
    provider_name?: string;
}

import { DefaultService } from '@/lib/api';

// ...

// ModelsResponse is the structure inside BaseResponse.data
interface ModelsResponse {
    image: Record<string, ModelConfig>;
    video: Record<string, ModelConfig>;
    audio: Record<string, ModelConfig>;
    lyrics: Record<string, ModelConfig>;
    music: Record<string, ModelConfig>;
    llm: Record<string, ModelConfig>;
}

interface ModelStore {
    imageModels: Record<string, ModelConfig>;
    videoModels: Record<string, ModelConfig>;
    llmModels: Record<string, ModelConfig>;
    audioModels: Record<string, ModelConfig>;
    lyricsModels: Record<string, ModelConfig>;
    musicModels: Record<string, ModelConfig>;
    defaults: {
        image: string | null;
        video: string | null;
        llm: string | null;
        audio: string | null;
        lyrics: string | null;
        music: string | null;
    };
    providers: any[]; // Cache for providers
    isLoading: boolean;
    error: string | null;
    setModels: (response: ModelsResponse) => void;
    fetchDefaults: () => Promise<void>;
    fetchProviders: () => Promise<void>;
    getDefaultModelId: (type: 'image' | 'video' | 'llm' | 'audio' | 'lyrics' | 'music') => string;
}

export const useModelStore = create<ModelStore>((set, get) => ({
    imageModels: {},
    videoModels: {},
    llmModels: {},
    audioModels: {},
    lyricsModels: {},
    musicModels: {},
    defaults: {
        image: null,
        video: null,
        llm: null,
        audio: null,
        lyrics: null,
        music: null
    },
    providers: [],
    isLoading: false,
    error: null,
    setModels: (response: ModelsResponse) => {
        // Backend now returns grouped HashMaps by type
        // Each type contains a HashMap with composite ID as key
        const imageModels: Record<string, ModelConfig> = {};
        const videoModels: Record<string, ModelConfig> = {};
        const llmModels: Record<string, ModelConfig> = {};
        const audioModels: Record<string, ModelConfig> = {};
        const lyricsModels: Record<string, ModelConfig> = {};
        const musicModels: Record<string, ModelConfig> = {};

        // Process image models
        Object.entries(response.image || {}).forEach(([id, model]) => {
            if (model.enabled === false) return;
            const icon = ICON_MAP[model.provider] || Qwen;
            imageModels[id] = { 
                ...model, 
                icon,
                capabilities: model.capabilities || {}
            };
        });

        // Process video models
        Object.entries(response.video || {}).forEach(([id, model]) => {
            if (model.enabled === false) return;
            const icon = ICON_MAP[model.provider] || Qwen;
            videoModels[id] = { 
                ...model, 
                icon,
                capabilities: model.capabilities || {}
            };
        });

        // Process LLM models
        Object.entries(response.llm || {}).forEach(([id, model]) => {
            if (model.enabled === false) return;
            const icon = ICON_MAP[model.provider] || Qwen;
            llmModels[id] = { 
                ...model, 
                icon,
                capabilities: model.capabilities || {}
            };
        });

        // Process audio models
        Object.entries(response.audio || {}).forEach(([id, model]) => {
            if (model.enabled === false) return;
            const icon = ICON_MAP[model.provider] || Qwen;
            audioModels[id] = { 
                ...model, 
                icon,
                capabilities: model.capabilities || {}
            };
        });

        // Process lyrics models
        Object.entries(response.lyrics || {}).forEach(([id, model]) => {
            if (model.enabled === false) return;
            const icon = ICON_MAP[model.provider] || Qwen;
            lyricsModels[id] = {
                ...model,
                icon,
                capabilities: model.capabilities || {}
            };
        });

        // Process music models
        Object.entries(response.music || {}).forEach(([id, model]) => {
            if (model.enabled === false) return;
            const icon = ICON_MAP[model.provider] || Qwen;
            musicModels[id] = {
                ...model,
                icon,
                capabilities: model.capabilities || {}
            };
        });

        set({ imageModels, videoModels, llmModels, audioModels, lyricsModels, musicModels, isLoading: false });
    },
    fetchDefaults: async () => {
        try {
            const response = await DefaultService.getDefaults();
            const data = response?.data as { image?: string; video?: string; llm?: string; audio?: string; lyrics?: string; music?: string; } | undefined;
            if (data) {
                set({ defaults: {
                    image: data.image || null,
                    video: data.video || null,
                    llm: data.llm || null,
                    audio: data.audio || null,
                    lyrics: data.lyrics || null,
                    music: data.music || null,
                }});
            }
        } catch (error) {
            console.error('Failed to fetch default models:', error);
        }
    },
    fetchProviders: async () => {
        const { providers } = get();
        if (providers.length > 0) return; // Return if already cached

        try {
            const response = await DefaultService.getProviders();
            // Backend returns BaseResponse with data containing { providers: [...] }
            const providersList = response?.data?.providers || response?.data || [];
            set({ providers: providersList });
        } catch (error) {
            console.error('Failed to fetch providers:', error);
        }
    },
    getDefaultModelId: (type) => {
        const models = type === 'image' ? get().imageModels 
                     : type === 'video' ? get().videoModels
                     : type === 'llm' ? get().llmModels
                     : type === 'audio' ? get().audioModels
                     : type === 'lyrics' ? get().lyricsModels
                     : get().musicModels;
        
        const modelsArray = Object.values(models);
        
        // 1. Priority: Use model marked as is_default
        const markedDefault = modelsArray.find(m => m.is_default);
        if (markedDefault) return markedDefault.id;

        // 2. Priority: Use system defaults fetched from API
        const systemDefaultId = get().defaults[type];
        if (systemDefaultId && models[systemDefaultId]) {
            return systemDefaultId;
        }

        // 3. Priority: Use hardcoded defaults (Fallback)s as last resort
        const hardcodedDefaults = {
            image: 'dashscope-z-image',
            video: 'dashscope-wan2.6-video',
            llm: 'dashscope-qwen-plus',
            audio: 'volcengine-doubao-tts',
            lyrics: 'minimax/lyrics-2.5',
            music: 'minimax/music-2.5'
        };
        
        return hardcodedDefaults[type];
    }
}));
