use std::collections::BTreeMap;

use chrono::Utc;
use http::StatusCode;
use serde_json::json;
use uuid::Uuid;

use fakecloud_core::service::{AwsRequest, AwsResponse, AwsServiceError};

use crate::state::{
    AccessTokenData, CognitoState, CustomDomainConfig, Device, RefreshTokenData, UserImportJob,
    UserPoolDomain,
};

use super::{
    device_to_json, domain_description_to_json, ensure_user_pool_exists, import_job_to_json,
    require_str, validate_string_length, CognitoService,
};

impl CognitoService {
    pub(super) fn create_user_pool_domain(
        &self,
        req: &AwsRequest,
    ) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();

        let pool_id = require_str(&body, "UserPoolId")?;
        let domain = require_str(&body, "Domain")?;

        let custom_domain_config =
            body["CustomDomainConfig"]["CertificateArn"]
                .as_str()
                .map(|arn| CustomDomainConfig {
                    certificate_arn: arn.to_string(),
                });

        let mut accounts = self.state.write();
        let state = accounts.get_or_create(&req.account_id);

        ensure_user_pool_exists(state, pool_id)?;

        if state.domains.contains_key(domain) {
            return Err(AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "InvalidParameterException",
                format!("Domain {domain} is already associated with a user pool."),
            ));
        }

        let domain_obj = UserPoolDomain {
            user_pool_id: pool_id.to_string(),
            domain: domain.to_string(),
            status: "ACTIVE".to_string(),
            custom_domain_config,
            creation_date: Utc::now(),
        };

        state.domains.insert(domain.to_string(), domain_obj);

        Ok(AwsResponse::ok_json(json!({})))
    }

    pub(super) fn describe_user_pool_domain(
        &self,
        req: &AwsRequest,
    ) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();

        let domain = require_str(&body, "Domain")?;
        validate_string_length(domain, "domain", 1, 63)?;

        let accounts = self.state.read();
        let empty = CognitoState::new(&req.account_id, &req.region);
        let state = accounts.get(&req.account_id).unwrap_or(&empty);

        // AWS returns empty DomainDescription if not found (no error)
        let description = match state.domains.get(domain) {
            Some(d) => domain_description_to_json(d, &state.account_id),
            None => json!({}),
        };

        Ok(AwsResponse::ok_json(json!({
            "DomainDescription": description
        })))
    }

    pub(super) fn update_user_pool_domain(
        &self,
        req: &AwsRequest,
    ) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();

        let pool_id = require_str(&body, "UserPoolId")?;
        let domain = require_str(&body, "Domain")?;

        let custom_domain_config =
            body["CustomDomainConfig"]["CertificateArn"]
                .as_str()
                .map(|arn| CustomDomainConfig {
                    certificate_arn: arn.to_string(),
                });

        let mut accounts = self.state.write();
        let state = accounts.get_or_create(&req.account_id);

        let d = state.domains.get_mut(domain).ok_or_else(|| {
            AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "ResourceNotFoundException",
                format!("Domain {domain} does not exist."),
            )
        })?;

        if d.user_pool_id != pool_id {
            return Err(AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "ResourceNotFoundException",
                format!("Domain {domain} does not exist."),
            ));
        }

        d.custom_domain_config = custom_domain_config;

        Ok(AwsResponse::ok_json(json!({})))
    }

    pub(super) fn delete_user_pool_domain(
        &self,
        req: &AwsRequest,
    ) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();

        let pool_id = require_str(&body, "UserPoolId")?;
        let domain = require_str(&body, "Domain")?;

        let mut accounts = self.state.write();
        let state = accounts.get_or_create(&req.account_id);

        match state.domains.get(domain) {
            Some(d) if d.user_pool_id == pool_id => {
                state.domains.remove(domain);
            }
            _ => {
                return Err(AwsServiceError::aws_error(
                    StatusCode::BAD_REQUEST,
                    "ResourceNotFoundException",
                    format!("Domain {domain} does not exist."),
                ));
            }
        }

        Ok(AwsResponse::ok_json(json!({})))
    }

    // ── Device Management ───────────────────────────────────────────────

    pub(super) fn admin_get_device(
        &self,
        req: &AwsRequest,
    ) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();
        let pool_id = require_str(&body, "UserPoolId")?;
        let username = require_str(&body, "Username")?;
        let device_key = require_str(&body, "DeviceKey")?;

        let accounts = self.state.read();
        let empty = CognitoState::new(&req.account_id, &req.region);
        let state = accounts.get(&req.account_id).unwrap_or(&empty);

        let users = state.users.get(pool_id).ok_or_else(|| {
            AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "ResourceNotFoundException",
                format!("User pool {pool_id} does not exist."),
            )
        })?;

        let user = users.get(username).ok_or_else(|| {
            AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "ResourceNotFoundException",
                format!("User {username} does not exist."),
            )
        })?;

        let device = user.devices.get(device_key).ok_or_else(|| {
            AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "ResourceNotFoundException",
                format!("Device {device_key} does not exist."),
            )
        })?;

        Ok(AwsResponse::ok_json(json!({
            "Device": device_to_json(device)
        })))
    }

    pub(super) fn admin_list_devices(
        &self,
        req: &AwsRequest,
    ) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();
        let pool_id = require_str(&body, "UserPoolId")?;
        let username = require_str(&body, "Username")?;
        let limit = body["Limit"].as_i64().unwrap_or(10).clamp(1, 60) as usize;
        let pagination_token = body["PaginationToken"].as_str();

        let accounts = self.state.read();
        let empty = CognitoState::new(&req.account_id, &req.region);
        let state = accounts.get(&req.account_id).unwrap_or(&empty);

        let users = state.users.get(pool_id).ok_or_else(|| {
            AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "ResourceNotFoundException",
                format!("User pool {pool_id} does not exist."),
            )
        })?;

        let user = users.get(username).ok_or_else(|| {
            AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "ResourceNotFoundException",
                format!("User {username} does not exist."),
            )
        })?;

        let mut devices: Vec<&Device> = user.devices.values().collect();
        devices.sort_by_key(|a| a.device_create_date);

        // A stale token (its device was forgotten) ends the listing rather
        // than silently restarting at page 1 (bug-audit 2026-06-20, 1.14).
        let start = match pagination_token {
            None => 0,
            Some(t) => devices
                .iter()
                .position(|d| d.device_key == t)
                .unwrap_or(devices.len()),
        };

        let page = &devices[start..devices.len().min(start + limit)];
        let next_token = if start + limit < devices.len() {
            devices.get(start + limit).map(|d| d.device_key.clone())
        } else {
            None
        };

        let mut result = json!({
            "Devices": page.iter().map(|d| device_to_json(d)).collect::<Vec<_>>()
        });
        if let Some(token) = next_token {
            result["PaginationToken"] = json!(token);
        }

        Ok(AwsResponse::ok_json(result))
    }

    pub(super) fn admin_forget_device(
        &self,
        req: &AwsRequest,
    ) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();
        let pool_id = require_str(&body, "UserPoolId")?;
        let username = require_str(&body, "Username")?;
        let device_key = require_str(&body, "DeviceKey")?;

        let mut accounts = self.state.write();
        let state = accounts.get_or_create(&req.account_id);

        let users = state.users.get_mut(pool_id).ok_or_else(|| {
            AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "ResourceNotFoundException",
                format!("User pool {pool_id} does not exist."),
            )
        })?;

        let user = users.get_mut(username).ok_or_else(|| {
            AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "ResourceNotFoundException",
                format!("User {username} does not exist."),
            )
        })?;

        if user.devices.remove(device_key).is_none() {
            return Err(AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "ResourceNotFoundException",
                format!("Device {device_key} does not exist."),
            ));
        }

        Ok(AwsResponse::ok_json(json!({})))
    }

    pub(super) fn admin_update_device_status(
        &self,
        req: &AwsRequest,
    ) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();
        let pool_id = require_str(&body, "UserPoolId")?;
        let username = require_str(&body, "Username")?;
        let device_key = require_str(&body, "DeviceKey")?;
        let status = body["DeviceRememberedStatus"]
            .as_str()
            .map(|s| s.to_string());

        let mut accounts = self.state.write();
        let state = accounts.get_or_create(&req.account_id);

        let users = state.users.get_mut(pool_id).ok_or_else(|| {
            AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "ResourceNotFoundException",
                format!("User pool {pool_id} does not exist."),
            )
        })?;

        let user = users.get_mut(username).ok_or_else(|| {
            AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "ResourceNotFoundException",
                format!("User {username} does not exist."),
            )
        })?;

        let device = user.devices.get_mut(device_key).ok_or_else(|| {
            AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "ResourceNotFoundException",
                format!("Device {device_key} does not exist."),
            )
        })?;

        device.device_remembered_status = status;
        device.device_last_modified_date = Utc::now();

        Ok(AwsResponse::ok_json(json!({})))
    }

    pub(super) fn confirm_device(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();
        let access_token = require_str(&body, "AccessToken")?;
        let device_key = require_str(&body, "DeviceKey")?;
        let device_name = body["DeviceName"].as_str();

        let mut accounts = self.state.write();
        let state = accounts.get_or_create(&req.account_id);

        let token_data = state.valid_access_token(access_token).ok_or_else(|| {
            AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "NotAuthorizedException",
                "Invalid access token.",
            )
        })?;
        let pool_id = token_data.user_pool_id.clone();
        let username = token_data.username.clone();

        let user = state
            .users
            .get_mut(&pool_id)
            .and_then(|users| users.get_mut(&username))
            .ok_or_else(|| {
                AwsServiceError::aws_error(
                    StatusCode::BAD_REQUEST,
                    "NotAuthorizedException",
                    "Invalid access token.",
                )
            })?;

        let now = Utc::now();
        let mut device_attributes = BTreeMap::new();
        if let Some(name) = device_name {
            device_attributes.insert("device_name".to_string(), name.to_string());
        }

        user.devices.insert(
            device_key.to_string(),
            Device {
                device_key: device_key.to_string(),
                device_attributes,
                device_create_date: now,
                device_last_modified_date: now,
                device_last_authenticated_date: Some(now),
                device_remembered_status: None,
            },
        );

        Ok(AwsResponse::ok_json(json!({
            "UserConfirmationNecessary": false
        })))
    }

    // ── User-facing Device Operations ─────────────────────────────────

    pub(super) fn forget_device(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();
        let access_token = body["AccessToken"].as_str();
        let device_key = require_str(&body, "DeviceKey")?;

        let mut accounts = self.state.write();
        let state = accounts.get_or_create(&req.account_id);

        let (pool_id, username) = if let Some(token) = access_token {
            let td = state.valid_access_token(token).ok_or_else(|| {
                AwsServiceError::aws_error(
                    StatusCode::BAD_REQUEST,
                    "NotAuthorizedException",
                    "Invalid access token.",
                )
            })?;
            (td.user_pool_id.clone(), td.username.clone())
        } else {
            return Err(AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "InvalidParameterException",
                "AccessToken is required",
            ));
        };

        let user = state
            .users
            .get_mut(&pool_id)
            .and_then(|users| users.get_mut(&username))
            .ok_or_else(|| {
                AwsServiceError::aws_error(
                    StatusCode::BAD_REQUEST,
                    "ResourceNotFoundException",
                    "User not found.",
                )
            })?;

        if user.devices.remove(device_key).is_none() {
            return Err(AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "ResourceNotFoundException",
                format!("Device {device_key} does not exist."),
            ));
        }

        Ok(AwsResponse::ok_json(json!({})))
    }

    pub(super) fn get_device(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();
        let access_token = body["AccessToken"].as_str();
        let device_key = require_str(&body, "DeviceKey")?;

        let accounts = self.state.read();
        let empty = CognitoState::new(&req.account_id, &req.region);
        let state = accounts.get(&req.account_id).unwrap_or(&empty);

        let (pool_id, username) = if let Some(token) = access_token {
            let td = state.valid_access_token(token).ok_or_else(|| {
                AwsServiceError::aws_error(
                    StatusCode::BAD_REQUEST,
                    "NotAuthorizedException",
                    "Invalid access token.",
                )
            })?;
            (td.user_pool_id.clone(), td.username.clone())
        } else {
            return Err(AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "InvalidParameterException",
                "AccessToken is required",
            ));
        };

        let user = state
            .users
            .get(&pool_id)
            .and_then(|users| users.get(&username))
            .ok_or_else(|| {
                AwsServiceError::aws_error(
                    StatusCode::BAD_REQUEST,
                    "ResourceNotFoundException",
                    "User not found.",
                )
            })?;

        let device = user.devices.get(device_key).ok_or_else(|| {
            AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "ResourceNotFoundException",
                format!("Device {device_key} does not exist."),
            )
        })?;

        Ok(AwsResponse::ok_json(json!({
            "Device": device_to_json(device)
        })))
    }

    pub(super) fn list_devices(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();
        let access_token = require_str(&body, "AccessToken")?;
        let limit = body["Limit"].as_i64().unwrap_or(10).clamp(1, 60) as usize;
        let pagination_token = body["PaginationToken"].as_str();

        let accounts = self.state.read();
        let empty = CognitoState::new(&req.account_id, &req.region);
        let state = accounts.get(&req.account_id).unwrap_or(&empty);

        let token_data = state.valid_access_token(access_token).ok_or_else(|| {
            AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "NotAuthorizedException",
                "Invalid access token.",
            )
        })?;
        let pool_id = &token_data.user_pool_id;
        let username = &token_data.username;

        let user = state
            .users
            .get(pool_id.as_str())
            .and_then(|users| users.get(username.as_str()))
            .ok_or_else(|| {
                AwsServiceError::aws_error(
                    StatusCode::BAD_REQUEST,
                    "ResourceNotFoundException",
                    "User not found.",
                )
            })?;

        let mut devices: Vec<&Device> = user.devices.values().collect();
        devices.sort_by_key(|a| a.device_create_date);

        // A stale token (its device was forgotten) ends the listing rather
        // than silently restarting at page 1 (bug-audit 2026-06-20, 1.14).
        let start = match pagination_token {
            None => 0,
            Some(t) => devices
                .iter()
                .position(|d| d.device_key == t)
                .unwrap_or(devices.len()),
        };

        let page = &devices[start..devices.len().min(start + limit)];
        let next_token = if start + limit < devices.len() {
            devices.get(start + limit).map(|d| d.device_key.clone())
        } else {
            None
        };

        let mut result = json!({
            "Devices": page.iter().map(|d| device_to_json(d)).collect::<Vec<_>>()
        });
        if let Some(token) = next_token {
            result["PaginationToken"] = json!(token);
        }

        Ok(AwsResponse::ok_json(result))
    }

    pub(super) fn update_device_status(
        &self,
        req: &AwsRequest,
    ) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();
        let access_token = require_str(&body, "AccessToken")?;
        let device_key = require_str(&body, "DeviceKey")?;
        let status = body["DeviceRememberedStatus"]
            .as_str()
            .map(|s| s.to_string());

        let mut accounts = self.state.write();
        let state = accounts.get_or_create(&req.account_id);

        let token_data = state.valid_access_token(access_token).ok_or_else(|| {
            AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "NotAuthorizedException",
                "Invalid access token.",
            )
        })?;
        let pool_id = token_data.user_pool_id.clone();
        let username = token_data.username.clone();

        let user = state
            .users
            .get_mut(&pool_id)
            .and_then(|users| users.get_mut(&username))
            .ok_or_else(|| {
                AwsServiceError::aws_error(
                    StatusCode::BAD_REQUEST,
                    "ResourceNotFoundException",
                    "User not found.",
                )
            })?;

        let device = user.devices.get_mut(device_key).ok_or_else(|| {
            AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "ResourceNotFoundException",
                format!("Device {device_key} does not exist."),
            )
        })?;

        device.device_remembered_status = status;
        device.device_last_modified_date = Utc::now();

        Ok(AwsResponse::ok_json(json!({})))
    }

    // ── Token Operations ───────────────────────────────────────────────

    pub(super) fn revoke_token(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();
        let token = require_str(&body, "Token")?;
        let client_id = require_str(&body, "ClientId")?;
        // ClientSecret is optional
        let _client_secret = body["ClientSecret"].as_str();

        let mut accounts = self.state.write();
        let state = accounts.get_or_create(&req.account_id);

        // Validate client exists
        if !state.user_pool_clients.contains_key(client_id) {
            return Err(AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "InvalidParameterException",
                format!("Client {client_id} does not exist."),
            ));
        }

        // Remove the refresh token if it exists
        state.refresh_tokens.remove(token);

        Ok(AwsResponse::ok_json(json!({})))
    }

    pub(super) fn get_tokens_from_refresh_token(
        &self,
        req: &AwsRequest,
    ) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();
        let refresh_token = require_str(&body, "RefreshToken")?;
        let client_id = require_str(&body, "ClientId")?;
        // Optional fields
        let _client_secret = body["ClientSecret"].as_str();
        let _device_key = body["DeviceKey"].as_str();

        let mut accounts = self.state.write();
        let state = accounts.get_or_create(&req.account_id);

        // Validate client exists
        if !state.user_pool_clients.contains_key(client_id) {
            return Err(AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "InvalidParameterException",
                format!("Client {client_id} does not exist."),
            ));
        }

        let token_data = state.refresh_tokens.get(refresh_token).ok_or_else(|| {
            AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "NotAuthorizedException",
                "Invalid refresh token.",
            )
        })?;

        // Verify client matches
        if token_data.client_id != client_id {
            return Err(AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "NotAuthorizedException",
                "Invalid refresh token.",
            ));
        }

        let pool_id = token_data.user_pool_id.clone();
        let username = token_data.username.clone();

        // Find user to get sub
        let user = state
            .users
            .get(pool_id.as_str())
            .and_then(|users| users.get(username.as_str()))
            .ok_or_else(|| {
                AwsServiceError::aws_error(
                    StatusCode::BAD_REQUEST,
                    "NotAuthorizedException",
                    "User not found.",
                )
            })?;

        let sub = user.sub.clone();
        let region = state.region.clone();

        let pool_signing_owned = state.user_pools.get(&pool_id).and_then(|pool| {
            pool.signing_key_pem
                .as_ref()
                .zip(pool.signing_kid.as_ref())
                .map(|(p, k)| (p.clone(), k.clone()))
        });
        let signing = pool_signing_owned
            .as_ref()
            .map(|(p, k)| (p.as_str(), k.as_str()));
        let tokens = super::generate_tokens(&pool_id, client_id, &sub, &username, &region, signing);

        let rotation_enabled = state
            .user_pool_clients
            .get(client_id)
            .and_then(|c| c.refresh_token_rotation.as_ref())
            .map(|r| r.feature.eq_ignore_ascii_case("ENABLED"))
            .unwrap_or(false);

        let now = Utc::now();
        let rotated_refresh = if rotation_enabled {
            let new_token = format!("rt-{}", Uuid::new_v4());
            state.refresh_tokens.insert(
                new_token.clone(),
                RefreshTokenData {
                    user_pool_id: pool_id.clone(),
                    username: username.clone(),
                    client_id: client_id.to_string(),
                    issued_at: now,
                },
            );
            state.refresh_tokens.remove(refresh_token);
            Some(new_token)
        } else {
            None
        };

        state.access_tokens.insert(
            tokens.access_token.clone(),
            AccessTokenData {
                user_pool_id: pool_id,
                username,
                client_id: client_id.to_string(),
                issued_at: now,
            },
        );

        let mut auth = json!({
            "AccessToken": tokens.access_token,
            "IdToken": tokens.id_token,
            "TokenType": "Bearer",
            "ExpiresIn": 3600
        });
        if let Some(ref rt) = rotated_refresh {
            auth["RefreshToken"] = json!(rt);
        }
        Ok(AwsResponse::ok_json(json!({
            "AuthenticationResult": auth,
        })))
    }

    // ── Tags ────────────────────────────────────────────────────────────

    pub(super) fn tag_resource(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();
        let resource_arn = require_str(&body, "ResourceArn")?;

        let tags: BTreeMap<String, String> = body["Tags"]
            .as_object()
            .map(|m| {
                m.iter()
                    .filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_string())))
                    .collect()
            })
            .unwrap_or_default();

        let mut accounts = self.state.write();
        let state = accounts.get_or_create(&req.account_id);

        // Validate that the ARN matches a known user pool
        let pool_exists = state.user_pools.values().any(|p| p.arn == resource_arn);
        if !pool_exists {
            return Err(AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "ResourceNotFoundException",
                format!("Resource {resource_arn} not found."),
            ));
        }

        let existing = state.tags.entry(resource_arn.to_string()).or_default();
        for (k, v) in tags {
            existing.insert(k, v);
        }

        Ok(AwsResponse::ok_json(json!({})))
    }

    pub(super) fn untag_resource(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();
        let resource_arn = require_str(&body, "ResourceArn")?;

        let tag_keys: Vec<String> = body["TagKeys"]
            .as_array()
            .map(|arr| {
                arr.iter()
                    .filter_map(|v| v.as_str().map(String::from))
                    .collect()
            })
            .unwrap_or_default();

        let mut accounts = self.state.write();
        let state = accounts.get_or_create(&req.account_id);

        let pool_exists = state.user_pools.values().any(|p| p.arn == resource_arn);
        if !pool_exists {
            return Err(AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "ResourceNotFoundException",
                format!("Resource {resource_arn} not found."),
            ));
        }

        if let Some(tags) = state.tags.get_mut(resource_arn) {
            for key in &tag_keys {
                tags.remove(key);
            }
        }

        Ok(AwsResponse::ok_json(json!({})))
    }

    pub(super) fn list_tags_for_resource(
        &self,
        req: &AwsRequest,
    ) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();
        let resource_arn = require_str(&body, "ResourceArn")?;

        let accounts = self.state.read();
        let empty = CognitoState::new(&req.account_id, &req.region);
        let state = accounts.get(&req.account_id).unwrap_or(&empty);

        let pool_exists = state.user_pools.values().any(|p| p.arn == resource_arn);
        if !pool_exists {
            return Err(AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "ResourceNotFoundException",
                format!("Resource {resource_arn} not found."),
            ));
        }

        let tags = state.tags.get(resource_arn).cloned().unwrap_or_default();

        Ok(AwsResponse::ok_json(json!({ "Tags": tags })))
    }

    // ── Import Jobs ─────────────────────────────────────────────────────

    pub(super) fn get_csv_header(&self, req: &AwsRequest) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();
        let pool_id = require_str(&body, "UserPoolId")?;

        let accounts = self.state.read();
        let empty = CognitoState::new(&req.account_id, &req.region);
        let state = accounts.get(&req.account_id).unwrap_or(&empty);

        let pool = state.user_pools.get(pool_id).ok_or_else(|| {
            AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "ResourceNotFoundException",
                format!("User pool {pool_id} does not exist."),
            )
        })?;

        let csv_header: Vec<String> = pool
            .schema_attributes
            .iter()
            .map(|a| a.name.clone())
            .collect();

        Ok(AwsResponse::ok_json(json!({
            "UserPoolId": pool_id,
            "CSVHeader": csv_header
        })))
    }

    pub(super) fn create_user_import_job(
        &self,
        req: &AwsRequest,
    ) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();
        let pool_id = require_str(&body, "UserPoolId")?;
        let job_name = require_str(&body, "JobName")?;
        let cw_role_arn = require_str(&body, "CloudWatchLogsRoleArn")?;

        let mut accounts = self.state.write();
        let state = accounts.get_or_create(&req.account_id);

        ensure_user_pool_exists(state, pool_id)?;

        let job_id = format!("import-{}", Uuid::new_v4());
        let now = Utc::now();

        let job = UserImportJob {
            job_id: job_id.clone(),
            job_name: job_name.to_string(),
            user_pool_id: pool_id.to_string(),
            cloud_watch_logs_role_arn: cw_role_arn.to_string(),
            status: "Created".to_string(),
            creation_date: now,
            start_date: None,
            completion_date: None,
            pre_signed_url: Some(format!(
                "https://fakecloud-import.s3.amazonaws.com/{pool_id}/{job_id}/upload.csv"
            )),
        };

        let resp = import_job_to_json(&job);

        state
            .import_jobs
            .entry(pool_id.to_string())
            .or_default()
            .insert(job_id, job);

        Ok(AwsResponse::ok_json(json!({
            "UserImportJob": resp
        })))
    }

    pub(super) fn describe_user_import_job(
        &self,
        req: &AwsRequest,
    ) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();
        let pool_id = require_str(&body, "UserPoolId")?;
        let job_id = require_str(&body, "JobId")?;

        let accounts = self.state.read();
        let empty = CognitoState::new(&req.account_id, &req.region);
        let state = accounts.get(&req.account_id).unwrap_or(&empty);

        let job = state
            .import_jobs
            .get(pool_id)
            .and_then(|jobs| jobs.get(job_id))
            .ok_or_else(|| {
                AwsServiceError::aws_error(
                    StatusCode::BAD_REQUEST,
                    "ResourceNotFoundException",
                    format!("Import job {job_id} does not exist."),
                )
            })?;

        Ok(AwsResponse::ok_json(json!({
            "UserImportJob": import_job_to_json(job)
        })))
    }

    pub(super) fn list_user_import_jobs(
        &self,
        req: &AwsRequest,
    ) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();
        let pool_id = require_str(&body, "UserPoolId")?;
        let max_results = body["MaxResults"].as_i64().unwrap_or(10).clamp(1, 60) as usize;
        let pagination_token = body["PaginationToken"].as_str();

        let accounts = self.state.read();
        let empty = CognitoState::new(&req.account_id, &req.region);
        let state = accounts.get(&req.account_id).unwrap_or(&empty);

        ensure_user_pool_exists(state, pool_id)?;

        let mut jobs: Vec<&UserImportJob> = state
            .import_jobs
            .get(pool_id)
            .map(|m| m.values().collect())
            .unwrap_or_default();
        jobs.sort_by_key(|a| a.creation_date);

        // A stale token (its job was pruned) ends the listing rather than
        // silently restarting at page 1 (bug-audit 2026-06-20, 1.14).
        let start = match pagination_token {
            None => 0,
            Some(t) => jobs
                .iter()
                .position(|j| j.job_id == t)
                .unwrap_or(jobs.len()),
        };

        let page = &jobs[start..jobs.len().min(start + max_results)];
        let next_token = if start + max_results < jobs.len() {
            jobs.get(start + max_results).map(|j| j.job_id.clone())
        } else {
            None
        };

        let mut result = json!({
            "UserImportJobs": page.iter().map(|j| import_job_to_json(j)).collect::<Vec<_>>()
        });
        if let Some(token) = next_token {
            result["PaginationToken"] = json!(token);
        }

        Ok(AwsResponse::ok_json(result))
    }

    pub(super) fn start_user_import_job(
        &self,
        req: &AwsRequest,
    ) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();
        let pool_id = require_str(&body, "UserPoolId")?;
        let job_id = require_str(&body, "JobId")?;

        let mut accounts = self.state.write();
        let state = accounts.get_or_create(&req.account_id);

        let job = state
            .import_jobs
            .get_mut(pool_id)
            .and_then(|jobs| jobs.get_mut(job_id))
            .ok_or_else(|| {
                AwsServiceError::aws_error(
                    StatusCode::BAD_REQUEST,
                    "ResourceNotFoundException",
                    format!("Import job {job_id} does not exist."),
                )
            })?;

        // Can only start from Created or Stopped status
        if job.status != "Created" && job.status != "Stopped" {
            return Err(AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "InvalidParameterException",
                format!("Import job cannot be started from status {}.", job.status),
            ));
        }

        let now = Utc::now();
        job.status = "InProgress".to_string();
        job.start_date = Some(now);

        let resp = import_job_to_json(job);

        Ok(AwsResponse::ok_json(json!({
            "UserImportJob": resp
        })))
    }

    pub(super) fn stop_user_import_job(
        &self,
        req: &AwsRequest,
    ) -> Result<AwsResponse, AwsServiceError> {
        let body = req.json_body();
        let pool_id = require_str(&body, "UserPoolId")?;
        let job_id = require_str(&body, "JobId")?;

        let mut accounts = self.state.write();
        let state = accounts.get_or_create(&req.account_id);

        let job = state
            .import_jobs
            .get_mut(pool_id)
            .and_then(|jobs| jobs.get_mut(job_id))
            .ok_or_else(|| {
                AwsServiceError::aws_error(
                    StatusCode::BAD_REQUEST,
                    "ResourceNotFoundException",
                    format!("Import job {job_id} does not exist."),
                )
            })?;

        if job.status != "InProgress" {
            return Err(AwsServiceError::aws_error(
                StatusCode::BAD_REQUEST,
                "InvalidParameterException",
                format!("Import job cannot be stopped from status {}.", job.status),
            ));
        }

        let now = Utc::now();
        job.status = "Stopped".to_string();
        job.completion_date = Some(now);

        let resp = import_job_to_json(job);

        Ok(AwsResponse::ok_json(json!({
            "UserImportJob": resp
        })))
    }
}
