use async_trait::async_trait;
use aws_credential_types::{Credentials, provider::SharedCredentialsProvider};
use aws_sdk_s3::Client as S3Client;
use aws_sdk_s3::config::{Builder as S3ConfigBuilder, Region};
use aws_sdk_s3::primitives::ByteStream;
use base64::Engine;
use base64::engine::general_purpose::STANDARD;
use reqwest::Url;
use sha2::{Digest, Sha256};

use crate::config::{AuthConfig, AvatarStorageConfig, AvatarStorageUrlStyle, extract_origin};
use crate::error::{AuthError, AuthResult};

const DEFAULT_AVATAR_STORAGE_REGION: &str = "auto";
const DATA_URL_MARKER: &str = ";base64,";

/// Error message returned when an external avatar URL origin is not trusted.
pub const UNTRUSTED_AVATAR_URL_MESSAGE: &str = "Avatar URL is not allowed by TRUSTED_ORIGINS";

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AvatarTargetKind {
    User,
    Organization,
}

impl AvatarTargetKind {
    pub fn as_template_value(self) -> &'static str {
        match self {
            Self::User => "user",
            Self::Organization => "organization",
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AvatarTemplateContext<'a> {
    pub target_kind: AvatarTargetKind,
    pub entity_id: &'a str,
    pub user_id: Option<&'a str>,
    pub organization_id: Option<&'a str>,
}

impl<'a> AvatarTemplateContext<'a> {
    pub fn for_user(user_id: &'a str) -> Self {
        Self {
            target_kind: AvatarTargetKind::User,
            entity_id: user_id,
            user_id: Some(user_id),
            organization_id: None,
        }
    }

    pub fn for_organization(organization_id: &'a str) -> Self {
        Self {
            target_kind: AvatarTargetKind::Organization,
            entity_id: organization_id,
            user_id: None,
            organization_id: Some(organization_id),
        }
    }

    pub fn with_user_id(mut self, user_id: &'a str) -> Self {
        self.user_id = Some(user_id);
        self
    }

    pub fn with_organization_id(mut self, organization_id: &'a str) -> Self {
        self.organization_id = Some(organization_id);
        self
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct PreparedAvatarObject {
    key: String,
    public_url: String,
    content_type: &'static str,
    bytes: Vec<u8>,
}

#[async_trait]
trait AvatarObjectStore: Send + Sync {
    async fn put_object(
        &self,
        storage: &AvatarStorageConfig,
        object: &PreparedAvatarObject,
    ) -> AuthResult<()>;
}

struct S3AvatarObjectStore;

#[async_trait]
impl AvatarObjectStore for S3AvatarObjectStore {
    async fn put_object(
        &self,
        storage: &AvatarStorageConfig,
        object: &PreparedAvatarObject,
    ) -> AuthResult<()> {
        let credentials = Credentials::new(
            storage.access_key.clone(),
            storage.secret_key.clone(),
            None,
            None,
            "athena-auth-avatar-storage",
        );
        let sdk_config = aws_config::defaults(aws_config::BehaviorVersion::latest())
            .credentials_provider(SharedCredentialsProvider::new(credentials))
            .region(Region::new(DEFAULT_AVATAR_STORAGE_REGION))
            .load()
            .await;

        let s3_config = S3ConfigBuilder::from(&sdk_config)
            .endpoint_url(storage.endpoint.trim())
            .force_path_style(matches!(storage.url_style, AvatarStorageUrlStyle::Path))
            .build();
        let client = S3Client::from_conf(s3_config);

        client
            .put_object()
            .bucket(storage.bucket.trim())
            .key(&object.key)
            .content_type(object.content_type)
            .body(ByteStream::from(object.bytes.clone()))
            .send()
            .await
            .map_err(|error| {
                AuthError::internal(format!(
                    "Failed to upload avatar object to storage: {error}"
                ))
            })?;

        Ok(())
    }
}

pub async fn resolve_avatar_value(
    config: &AuthConfig,
    raw_value: &str,
    template_context: AvatarTemplateContext<'_>,
) -> AuthResult<String> {
    let store = S3AvatarObjectStore;
    resolve_avatar_value_with_store(config, raw_value, template_context, &store).await
}

async fn resolve_avatar_value_with_store(
    config: &AuthConfig,
    raw_value: &str,
    template_context: AvatarTemplateContext<'_>,
    store: &dyn AvatarObjectStore,
) -> AuthResult<String> {
    let trimmed = raw_value.trim();
    if trimmed.is_empty() {
        return Ok(raw_value.to_string());
    }

    if extract_origin(trimmed).is_some() {
        if config.is_origin_trusted(trimmed) {
            return Ok(trimmed.to_string());
        }

        return Err(AuthError::bad_request(UNTRUSTED_AVATAR_URL_MESSAGE));
    }

    if !config.avatar_storage.enabled {
        return Ok(raw_value.to_string());
    }

    let object = prepare_avatar_object(&config.avatar_storage, trimmed, template_context)?;
    store.put_object(&config.avatar_storage, &object).await?;
    Ok(object.public_url)
}

fn prepare_avatar_object(
    storage: &AvatarStorageConfig,
    raw_value: &str,
    template_context: AvatarTemplateContext<'_>,
) -> AuthResult<PreparedAvatarObject> {
    let decoded = decode_avatar_payload(raw_value)?;
    let detected = detect_avatar_image_format(&decoded)?;
    let sha256 = hex_sha256(&decoded);
    let key = render_avatar_object_key(storage, template_context, &sha256, detected.extension);
    let public_url = build_avatar_public_url(storage, &key)?;

    Ok(PreparedAvatarObject {
        key,
        public_url,
        content_type: detected.content_type,
        bytes: decoded,
    })
}

fn decode_avatar_payload(raw_value: &str) -> AuthResult<Vec<u8>> {
    if let Some(encoded) = raw_value.strip_prefix("data:") {
        let Some(marker_index) = encoded.find(DATA_URL_MARKER) else {
            return Err(AuthError::bad_request(
                "Avatar data URL must use base64 encoding",
            ));
        };
        let payload = &encoded[(marker_index + DATA_URL_MARKER.len())..];
        return STANDARD.decode(payload).map_err(|error| {
            AuthError::bad_request(format!("Invalid avatar base64 payload: {error}"))
        });
    }

    STANDARD
        .decode(raw_value)
        .map_err(|error| AuthError::bad_request(format!("Invalid avatar base64 payload: {error}")))
}

struct DetectedAvatarFormat {
    extension: &'static str,
    content_type: &'static str,
}

fn detect_avatar_image_format(bytes: &[u8]) -> AuthResult<DetectedAvatarFormat> {
    if bytes.starts_with(&[0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A]) {
        return Ok(DetectedAvatarFormat {
            extension: "png",
            content_type: "image/png",
        });
    }
    if bytes.starts_with(&[0xFF, 0xD8, 0xFF]) {
        return Ok(DetectedAvatarFormat {
            extension: "jpg",
            content_type: "image/jpeg",
        });
    }
    if bytes.len() >= 12 && &bytes[0..4] == b"RIFF" && &bytes[8..12] == b"WEBP" {
        return Ok(DetectedAvatarFormat {
            extension: "webp",
            content_type: "image/webp",
        });
    }

    Err(AuthError::bad_request(
        "Avatar image must be a PNG, JPEG, or WebP file",
    ))
}

fn render_avatar_object_key(
    storage: &AvatarStorageConfig,
    template_context: AvatarTemplateContext<'_>,
    sha256: &str,
    extension: &str,
) -> String {
    let rendered = render_avatar_template_expression(
        &storage.file_key_template,
        template_context,
        sha256,
        extension,
    );

    let prefix = storage
        .prefix_path
        .as_deref()
        .map(str::trim)
        .unwrap_or_default()
        .trim_matches('/');
    let rendered = rendered.trim().trim_matches('/');

    match (prefix.is_empty(), rendered.is_empty()) {
        (true, true) => String::new(),
        (true, false) => rendered.to_string(),
        (false, true) => prefix.to_string(),
        (false, false) => format!("{prefix}/{rendered}"),
    }
}

fn render_avatar_template_expression(
    template: &str,
    template_context: AvatarTemplateContext<'_>,
    sha256: &str,
    extension: &str,
) -> String {
    let mut rendered = String::with_capacity(template.len());
    let mut cursor = template;

    while let Some(start) = cursor.find('{') {
        rendered.push_str(&cursor[..start]);
        let remainder = &cursor[(start + 1)..];
        let Some(end) = remainder.find('}') else {
            rendered.push('{');
            rendered.push_str(remainder);
            return rendered;
        };

        let expression = &remainder[..end];
        rendered.push_str(resolve_avatar_template_expression(
            expression,
            template_context,
            sha256,
            extension,
        ));
        cursor = &remainder[(end + 1)..];
    }

    rendered.push_str(cursor);
    rendered
}

fn resolve_avatar_template_expression<'a>(
    expression: &'a str,
    template_context: AvatarTemplateContext<'a>,
    sha256: &'a str,
    extension: &'a str,
) -> &'a str {
    for candidate in expression.split("||") {
        let Some(value) =
            resolve_avatar_template_variable(candidate.trim(), template_context, sha256, extension)
        else {
            continue;
        };
        if !value.is_empty() {
            return value;
        }
    }

    ""
}

fn resolve_avatar_template_variable<'a>(
    name: &str,
    template_context: AvatarTemplateContext<'a>,
    sha256: &'a str,
    extension: &'a str,
) -> Option<&'a str> {
    match name {
        "kind" => Some(template_context.target_kind.as_template_value()),
        "entity_id" => Some(template_context.entity_id),
        "user_id" => Some(template_context.user_id.unwrap_or_default()),
        "organization_id" => Some(template_context.organization_id.unwrap_or_default()),
        "sha256" => Some(sha256),
        "extension" => Some(extension),
        _ => None,
    }
}

fn build_avatar_public_url(storage: &AvatarStorageConfig, key: &str) -> AuthResult<String> {
    let mut url = Url::parse(storage.endpoint.trim()).map_err(|error| {
        AuthError::config(format!(
            "Avatar storage endpoint is not a valid URL: {error}"
        ))
    })?;

    let key_segments: Vec<&str> = key
        .split('/')
        .map(str::trim)
        .filter(|segment| !segment.is_empty())
        .collect();
    if key_segments.is_empty() {
        return Err(AuthError::config(
            "Avatar storage file_key_template resolved to an empty object key",
        ));
    }

    match storage.url_style {
        AvatarStorageUrlStyle::Path => {
            let mut segments = url.path_segments_mut().map_err(|_| {
                AuthError::config("Avatar storage endpoint cannot be used for path-style URLs")
            })?;
            segments.push(storage.bucket.trim());
            for segment in &key_segments {
                segments.push(segment);
            }
        }
        AvatarStorageUrlStyle::VirtualHost => {
            let host = url
                .host_str()
                .ok_or_else(|| AuthError::config("Avatar storage endpoint must include a host"))?;
            let bucket_host = format!("{}.{}", storage.bucket.trim(), host);
            url.set_host(Some(&bucket_host)).map_err(|_| {
                AuthError::config("Avatar storage endpoint host could not be combined with bucket")
            })?;

            let mut segments = url.path_segments_mut().map_err(|_| {
                AuthError::config(
                    "Avatar storage endpoint cannot be used for virtual-host-style URLs",
                )
            })?;
            for segment in &key_segments {
                segments.push(segment);
            }
        }
    }

    url.set_query(None);
    url.set_fragment(None);
    Ok(url.to_string())
}

fn hex_sha256(bytes: &[u8]) -> String {
    let digest = Sha256::digest(bytes);
    let mut output = String::with_capacity(digest.len() * 2);
    for byte in digest {
        output.push_str(&format!("{byte:02x}"));
    }
    output
}

#[cfg(test)]
mod tests {
    use std::sync::Mutex;

    use super::*;
    use crate::config::AuthConfig;

    #[derive(Default)]
    struct RecordingStore {
        uploads: Mutex<Vec<PreparedAvatarObject>>,
    }

    #[async_trait]
    impl AvatarObjectStore for RecordingStore {
        async fn put_object(
            &self,
            _storage: &AvatarStorageConfig,
            object: &PreparedAvatarObject,
        ) -> AuthResult<()> {
            self.uploads
                .lock()
                .expect("lock should succeed")
                .push(object.clone());
            Ok(())
        }
    }

    fn avatar_enabled_config() -> AuthConfig {
        let mut config = AuthConfig::new("test-secret-key-that-is-at-least-32-characters-long")
            .trusted_origin("https://*.example.com");
        config.avatar_storage = AvatarStorageConfig {
            enabled: true,
            bucket: "avatars".to_string(),
            endpoint: "https://storage.example.com".to_string(),
            access_key: "key".to_string(),
            secret_key: "secret".to_string(),
            url_style: AvatarStorageUrlStyle::Path,
            prefix_path: Some("auth".to_string()),
            file_key_template: "{kind}/{entity_id}/{sha256}.{extension}".to_string(),
        };
        config
    }

    #[tokio::test]
    async fn trusted_url_is_preserved() {
        let config = avatar_enabled_config();
        let resolved = resolve_avatar_value_with_store(
            &config,
            "https://cdn.example.com/avatars/a.png",
            AvatarTemplateContext::for_user("user-1"),
            &RecordingStore::default(),
        )
        .await
        .expect("trusted URL should succeed");
        assert_eq!(resolved, "https://cdn.example.com/avatars/a.png");
    }

    #[tokio::test]
    async fn untrusted_url_is_rejected() {
        let config = avatar_enabled_config();
        let error = resolve_avatar_value_with_store(
            &config,
            "https://evil.example.net/avatars/a.png",
            AvatarTemplateContext::for_user("user-1"),
            &RecordingStore::default(),
        )
        .await
        .expect_err("untrusted URL should fail");
        assert!(error.to_string().contains("TRUSTED_ORIGINS"));
    }

    #[tokio::test]
    async fn storage_disabled_preserves_raw_base64() {
        let config = AuthConfig::new("test-secret-key-that-is-at-least-32-characters-long");
        let raw = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAusB9Wn8n2QAAAAASUVORK5CYII=";
        let resolved = resolve_avatar_value_with_store(
            &config,
            raw,
            AvatarTemplateContext::for_user("user-1"),
            &RecordingStore::default(),
        )
        .await
        .expect("raw base64 should pass through when storage is disabled");
        assert_eq!(resolved, raw);
    }

    #[tokio::test]
    async fn data_uri_uploads_and_returns_public_url() {
        let config = avatar_enabled_config();
        let store = RecordingStore::default();
        let resolved = resolve_avatar_value_with_store(
            &config,
            "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAusB9Wn8n2QAAAAASUVORK5CYII=",
            AvatarTemplateContext::for_organization("org-1"),
            &store,
        )
        .await
        .expect("data URI should upload");

        assert!(
            resolved.starts_with("https://storage.example.com/avatars/auth/organization/org-1/")
        );
        let uploads = store.uploads.lock().expect("lock should succeed");
        assert_eq!(uploads.len(), 1);
        assert_eq!(uploads[0].content_type, "image/png");
    }

    #[tokio::test]
    async fn malformed_base64_is_rejected() {
        let config = avatar_enabled_config();
        let error = resolve_avatar_value_with_store(
            &config,
            "%%%not-base64%%%",
            AvatarTemplateContext::for_user("user-1"),
            &RecordingStore::default(),
        )
        .await
        .expect_err("invalid base64 should fail");
        assert!(error.to_string().contains("Invalid avatar base64 payload"));
    }

    #[tokio::test]
    async fn unsupported_image_payload_is_rejected() {
        let config = avatar_enabled_config();
        let raw = STANDARD.encode(b"plain-text");
        let error = resolve_avatar_value_with_store(
            &config,
            &raw,
            AvatarTemplateContext::for_user("user-1"),
            &RecordingStore::default(),
        )
        .await
        .expect_err("unsupported payload should fail");
        assert!(
            error
                .to_string()
                .contains("Avatar image must be a PNG, JPEG, or WebP file")
        );
    }

    #[test]
    fn file_key_template_supports_user_and_organization_placeholders() {
        let storage = AvatarStorageConfig {
            enabled: true,
            bucket: "avatars".to_string(),
            endpoint: "https://storage.example.com".to_string(),
            access_key: "key".to_string(),
            secret_key: "secret".to_string(),
            url_style: AvatarStorageUrlStyle::Path,
            prefix_path: Some("auth".to_string()),
            file_key_template:
                "{kind}/{user_id}/{organization_id}/{entity_id}/{sha256}.{extension}".to_string(),
        };

        let rendered = render_avatar_object_key(
            &storage,
            AvatarTemplateContext::for_organization("org-1").with_user_id("user-1"),
            "hash",
            "png",
        );

        assert_eq!(rendered, "auth/organization/user-1/org-1/org-1/hash.png");
    }

    #[test]
    fn file_key_template_supports_fallback_expressions() {
        let storage = AvatarStorageConfig {
            enabled: true,
            bucket: "avatars".to_string(),
            endpoint: "https://storage.example.com".to_string(),
            access_key: "key".to_string(),
            secret_key: "secret".to_string(),
            url_style: AvatarStorageUrlStyle::Path,
            prefix_path: None,
            file_key_template:
                "/avatars/{organization_id || user_id}/{kind}/{entity_id}/{sha256}.{extension}"
                    .to_string(),
        };

        let organization_rendered = render_avatar_object_key(
            &storage,
            AvatarTemplateContext::for_organization("org-1").with_user_id("user-1"),
            "hash",
            "png",
        );
        let user_rendered = render_avatar_object_key(
            &storage,
            AvatarTemplateContext::for_user("user-1"),
            "hash",
            "png",
        );

        assert_eq!(
            organization_rendered,
            "avatars/org-1/organization/org-1/hash.png"
        );
        assert_eq!(user_rendered, "avatars/user-1/user/user-1/hash.png");
    }
}
