use std::collections::HashMap;
use std::sync::Arc;

use crate::api::OpenRouterImagesApi;
use crate::auth::{
    AuthContext, AuthModel, AuthResolutionOverrides, AuthResult, CredentialStore, InMemoryCredentialStore, ModelsError,
    ModelsErrorCode, ProviderAuth, ProviderAuthHolder, env_api_key_auth, resolve_provider_auth,
};
use crate::images::models::OPENROUTER_IMAGE_MODELS;
use crate::types::{AssistantImages, ImagesContext, ImagesModel, ImagesOptions, ProviderImages};

pub struct ImagesProvider {
    pub id: String,
    pub name: String,
    pub auth: ProviderAuth,
    models: Vec<ImagesModel>,
    api: Arc<dyn ProviderImages>,
}

impl ImagesProvider {
    pub fn get_models(&self) -> &[ImagesModel] {
        &self.models
    }
}

pub struct CreateImagesProviderOptions {
    pub id: String,
    pub name: Option<String>,
    pub auth: ProviderAuth,
    pub models: Vec<ImagesModel>,
    pub api: Arc<dyn ProviderImages>,
}

pub fn create_images_provider(input: CreateImagesProviderOptions) -> ImagesProvider {
    let id = input.id.clone();
    ImagesProvider {
        id: input.id,
        name: input.name.unwrap_or(id),
        auth: input.auth,
        models: input.models,
        api: input.api,
    }
}

pub struct CreateImagesModelsOptions {
    pub credentials: Option<Arc<dyn CredentialStore>>,
    pub auth_context: Option<Arc<dyn AuthContext>>,
}

pub struct ImagesModels {
    providers: HashMap<String, ImagesProvider>,
    credentials: Arc<dyn CredentialStore>,
    auth_context: Arc<dyn AuthContext>,
}

pub struct MutableImagesModels {
    inner: ImagesModels,
}

impl ImagesModels {
    pub fn get_providers(&self) -> Vec<&ImagesProvider> {
        self.providers.values().collect()
    }

    pub fn get_provider(&self, id: &str) -> Option<&ImagesProvider> {
        self.providers.get(id)
    }

    pub fn get_models(&self, provider: Option<&str>) -> Vec<ImagesModel> {
        match provider {
            Some(id) => self
                .providers
                .get(id)
                .map(|p| p.get_models().to_vec())
                .unwrap_or_default(),
            None => self
                .providers
                .values()
                .flat_map(|p| p.get_models().iter().cloned())
                .collect(),
        }
    }

    pub fn get_model(&self, provider: &str, id: &str) -> Option<ImagesModel> {
        self.get_models(Some(provider)).into_iter().find(|m| m.id == id)
    }

    pub async fn get_auth(&self, model: &ImagesModel) -> Result<Option<AuthResult>, ModelsError> {
        let provider = self.providers.get(&model.provider).ok_or_else(|| {
            ModelsError::new(
                ModelsErrorCode::Provider,
                format!("Unknown provider: {}", model.provider),
            )
        })?;
        resolve_provider_auth(
            &ProviderAuthHolder {
                id: provider.id.clone(),
                auth: provider.auth.clone(),
            },
            AuthModel::Images(model.clone()),
            self.credentials.as_ref(),
            self.auth_context.clone(),
            None,
        )
        .await
    }

    pub async fn generate_images(
        &self,
        model: &ImagesModel,
        context: &ImagesContext,
        options: Option<ImagesOptions>,
    ) -> AssistantImages {
        let Some(provider) = self.providers.get(&model.provider) else {
            return AssistantImages {
                api: model.api.clone(),
                provider: model.provider.clone(),
                model: model.id.clone(),
                output: vec![],
                response_id: None,
                usage: None,
                stop_reason: crate::types::StopReason::Error,
                error_message: Some(format!("Unknown provider: {}", model.provider)),
                timestamp: chrono::Utc::now().timestamp_millis(),
            };
        };

        let overrides = options.as_ref().map(|o| AuthResolutionOverrides {
            api_key: o.api_key.clone(),
            env: o.env.clone(),
        });
        let resolution = match resolve_provider_auth(
            &ProviderAuthHolder {
                id: provider.id.clone(),
                auth: provider.auth.clone(),
            },
            AuthModel::Images(model.clone()),
            self.credentials.as_ref(),
            self.auth_context.clone(),
            overrides,
        )
        .await
        {
            Ok(r) => r,
            Err(e) => {
                return AssistantImages {
                    api: model.api.clone(),
                    provider: model.provider.clone(),
                    model: model.id.clone(),
                    output: vec![],
                    response_id: None,
                    usage: None,
                    stop_reason: crate::types::StopReason::Error,
                    error_message: Some(e.message),
                    timestamp: chrono::Utc::now().timestamp_millis(),
                };
            }
        };

        let mut opts = options.unwrap_or(ImagesOptions {
            api_key: None,
            signal: None,
            env: None,
            headers: None,
            timeout_ms: None,
            max_retries: None,
            on_payload: None,
            on_response: None,
        });
        if let Some(res) = resolution {
            if opts.api_key.is_none() {
                opts.api_key = res.auth.api_key;
            }
            if let Some(headers) = res.auth.headers {
                let mut merged = headers;
                if let Some(request_headers) = opts.headers.take() {
                    merged.extend(request_headers);
                }
                opts.headers = Some(merged);
            }
            if res.env.is_some() || opts.env.is_some() {
                let mut merged = res.env.unwrap_or_default();
                if let Some(request_env) = opts.env.take() {
                    merged.extend(request_env);
                }
                opts.env = Some(merged);
            }
        }

        provider.api.generate_images(model, context, Some(opts)).await
    }
}

pub fn create_images_models(options: Option<CreateImagesModelsOptions>) -> MutableImagesModels {
    MutableImagesModels {
        inner: ImagesModels {
            providers: HashMap::new(),
            credentials: options
                .as_ref()
                .and_then(|o| o.credentials.clone())
                .unwrap_or_else(|| Arc::new(InMemoryCredentialStore::new())),
            auth_context: options
                .as_ref()
                .and_then(|o| o.auth_context.clone())
                .unwrap_or_else(|| Arc::new(crate::auth::DefaultAuthContext::new())),
        },
    }
}

impl MutableImagesModels {
    pub fn set_provider(&mut self, provider: ImagesProvider) {
        self.inner.providers.insert(provider.id.clone(), provider);
    }
}

impl std::ops::Deref for MutableImagesModels {
    type Target = ImagesModels;
    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

pub fn openrouter_images_provider() -> ImagesProvider {
    create_images_provider(CreateImagesProviderOptions {
        id: "openrouter".to_string(),
        name: Some("OpenRouter".to_string()),
        auth: ProviderAuth {
            api_key: Some(env_api_key_auth("OpenRouter API key", vec!["OPENROUTER_API_KEY"])),
            oauth: None,
        },
        models: OPENROUTER_IMAGE_MODELS.to_vec(),
        api: Arc::new(OpenRouterImagesApi),
    })
}

pub fn builtin_images_models(options: Option<CreateImagesModelsOptions>) -> MutableImagesModels {
    let mut models = create_images_models(options);
    models.set_provider(openrouter_images_provider());
    models
}

pub async fn generate_images(
    model: &ImagesModel,
    context: &ImagesContext,
    options: Option<ImagesOptions>,
) -> AssistantImages {
    let collection = builtin_images_models(None);
    collection.generate_images(model, context, options).await
}
