use alloy::primitives::Address;
use base64::{engine::general_purpose::STANDARD, Engine};
use newton_data_provider::wasm::secrets::SecretsProvider;
use newton_prover_core::database::{
    initialize_database, DatabaseConfig, DatabaseManager, PolicyClientOwnerRepository, WasmSecretRepository,
};
use tracing::info;
use uuid::Uuid;

/// Helper to set up test database with a test secret
async fn setup_test_database() -> Result<(DatabaseManager, Uuid, Address, Address), Box<dyn std::error::Error>> {
    // Create a test database connection
    let db_config = DatabaseConfig {
        url: std::env::var("DATABASE_URL")
            .unwrap_or_else(|_| "postgres://newton:newton@localhost:5432/newton_gateway".to_string()),
        max_connections: 5,
        min_connections: 1,
        connect_timeout: std::time::Duration::from_secs(10),
    };

    let db = DatabaseManager::new(db_config).await?;

    // Create test user
    let user_id = Uuid::new_v4();
    let policy_client_address = Address::from([0x42u8; 20]);
    let policy_data_address = Address::from([0x24u8; 20]);

    // Claim ownership for this policy client (required by FK).
    let owner_repo = PolicyClientOwnerRepository::new(db.clone());
    owner_repo
        .claim_if_unowned(user_id, policy_client_address)
        .await
        .map_err(|e| format!("failed to claim policy_client ownership: {e}"))?;

    // Insert test encrypted secrets blob into database.
    // NOTE: In production this would be KMS ciphertext bytes (BYTEA).
    // For tests we store arbitrary bytes (here: base64 string bytes of a JSON blob).
    let plaintext_json = r#"{"TEST_API_KEY":"mock_value"}"#;
    let secrets_blob = STANDARD.encode(plaintext_json.as_bytes()).into_bytes();

    sqlx::query(
        r#"
        INSERT INTO policy_client_secret (policy_client_address, policy_data_address, secrets)
        VALUES ($1, $2, $3)
        ON CONFLICT (policy_client_address, policy_data_address) DO UPDATE
        SET secrets = EXCLUDED.secrets
        "#,
    )
    .bind(policy_client_address.as_slice())
    .bind(policy_data_address.as_slice())
    .bind(secrets_blob)
    .execute(db.pool())
    .await?;

    Ok((db, user_id, policy_client_address, policy_data_address))
}

/// Helper to clean up test data
async fn cleanup_test_database(db: &DatabaseManager, user_id: Uuid) -> Result<(), Box<dyn std::error::Error>> {
    sqlx::query(
        r#"
        DELETE FROM policy_client_secret
        WHERE policy_client_address IN (
            SELECT policy_client_address FROM policy_client_owner WHERE user_id = $1
        )
        "#,
    )
    .bind(user_id)
    .execute(db.pool())
    .await?;

    sqlx::query(
        r#"
        DELETE FROM policy_client_owner WHERE user_id = $1
        "#,
    )
    .bind(user_id)
    .execute(db.pool())
    .await?;

    Ok(())
}

/// Tests that SecretsProvider can detect database availability and is ready
/// to retrieve secrets. This validates the availability check mechanism that
/// determines whether the secrets infrastructure is properly initialized.
///
/// Note: This test validates availability detection but doesn't test actual
/// KMS decryption, which would require AWS credentials in the test environment.
#[tokio::test]
async fn test_secrets_provider_retrieves_from_database() {
    // Try to set up test database
    let setup_result = setup_test_database().await;

    if setup_result.is_err() {
        info!(
            "Database not available or test table not created, skipping database integration test: {:?}",
            setup_result.err()
        );
        return;
    }

    let (db, user_id, policy_client_address, policy_data_address) = setup_result.unwrap();

    // Initialize the global database instance
    let db_config = DatabaseConfig {
        url: std::env::var("DATABASE_URL")
            .unwrap_or_else(|_| "postgres://newton:newton@localhost:5432/newton_gateway".to_string()),
        max_connections: 5,
        min_connections: 1,
        connect_timeout: std::time::Duration::from_secs(10),
    };

    // Try to initialize - it may already be initialized
    let _ = initialize_database(db_config).await;

    // Create SecretsProvider
    let secrets_provider = SecretsProvider::new(
        None,
        Some(policy_client_address),
        Some(policy_data_address),
        "test-kms-key".to_string(),
        None,
        None,
    );

    // Verify database is available
    assert!(
        secrets_provider.is_available(),
        "SecretsProvider should detect database is available"
    );

    // Note: Actually calling get_secret would require KMS credentials
    // In a real integration test environment with KMS, we would test:
    // let result = secrets_provider.get_secret("TEST_API_KEY").await;
    // assert!(result.is_ok());

    // Clean up
    cleanup_test_database(&db, user_id)
        .await
        .expect("Failed to cleanup test data");

    info!("SecretsProvider database integration test completed successfully");
}

/// Tests that WasmSecretRepository can successfully query and retrieve
/// encrypted secrets blobs from the database scoped by (policy_client_address, policy_data_address).
#[tokio::test]
async fn test_wasm_secret_repository_get_by_policy_data_and_policy_client() {
    let setup_result = setup_test_database().await;

    if setup_result.is_err() {
        info!(
            "Database not available, skipping WasmSecretRepository test: {:?}",
            setup_result.err()
        );
        return;
    }

    let (db, user_id, policy_client_address, policy_data_address) = setup_result.unwrap();

    let repo = WasmSecretRepository::new(db.clone());

    // Test get_by_policy_data_and_policy_client
    let result = repo
        .get_by_policy_data_and_policy_client(policy_data_address, policy_client_address)
        .await;

    assert!(result.is_ok(), "Failed to query secret by key: {:?}", result.err());

    let secret = result.unwrap();
    assert!(secret.is_some(), "Secret should be found");

    let secret_record = secret.unwrap();
    assert_eq!(secret_record.policy_client_address, policy_client_address);
    assert_eq!(secret_record.policy_data_address, policy_data_address);

    cleanup_test_database(&db, user_id).await.expect("Failed to cleanup");

    info!("WasmSecretRepository get_by_policy_data_and_policy_client test completed");
}

/// Tests policy client address scoping for secrets. Validates that:
/// 1. Secrets can be retrieved when the correct policy client address is provided
/// 2. Secrets are NOT accessible when using a different policy client address
///
/// This ensures multi-tenant secret isolation, preventing policy clients from
/// accessing each other's secrets.
#[tokio::test]
async fn test_wasm_secret_repository_get_by_policy_data_and_policy_client_scoped() {
    let setup_result = setup_test_database().await;

    if setup_result.is_err() {
        info!(
            "Database not available, skipping policy client scoped secret test: {:?}",
            setup_result.err()
        );
        return;
    }

    let (db, user_id, policy_client_address, policy_data_address) = setup_result.unwrap();

    let repo = WasmSecretRepository::new(db.clone());

    // Test get_by_policy_data_and_policy_client with correct address
    let result = repo
        .get_by_policy_data_and_policy_client(policy_data_address, policy_client_address)
        .await;

    assert!(
        result.is_ok(),
        "Failed to query secret by key and policy client: {:?}",
        result.err()
    );

    let secret = result.unwrap();
    assert!(secret.is_some(), "Secret should be found for correct policy client");

    // Test with wrong policy client address
    let wrong_address = Address::from([0x99u8; 20]);
    let result_wrong = repo
        .get_by_policy_data_and_policy_client(policy_data_address, wrong_address)
        .await;

    assert!(result_wrong.is_ok());
    assert!(
        result_wrong.unwrap().is_none(),
        "Secret should not be found for wrong policy client"
    );

    cleanup_test_database(&db, user_id).await.expect("Failed to cleanup");

    info!("WasmSecretRepository policy client scoped query test completed");
}

/// Tests error handling for non-existent secrets. Validates that the repository
/// gracefully handles queries for secrets that don't exist in the database,
/// returning None instead of erroring.
#[tokio::test]
async fn test_secrets_not_found() {
    let setup_result = setup_test_database().await;

    if setup_result.is_err() {
        info!(
            "Database not available, skipping non-existent secret test: {:?}",
            setup_result.err()
        );
        return;
    }

    let (db, user_id, policy_client_address, _) = setup_result.unwrap();

    let repo = WasmSecretRepository::new(db.clone());

    // Try to get a secrets blob that doesn't exist (wrong policy_data_address)
    let wrong_policy_data = Address::from([0x99u8; 20]);
    let result = repo
        .get_by_policy_data_and_policy_client(wrong_policy_data, policy_client_address)
        .await;

    assert!(result.is_ok());
    assert!(result.unwrap().is_none(), "Non-existent secret should return None");

    cleanup_test_database(&db, user_id).await.expect("Failed to cleanup");

    info!("Non-existent secret handling test completed");
}

/// Tests the SecretsProvider.is_available() method in both scenarios:
/// 1. Without database - should return false
/// 2. With database initialized - should return true
///
/// This availability check is critical for WASM modules to know whether they
/// can request secrets or need to handle missing secrets infrastructure.
#[tokio::test]
async fn test_secrets_provider_availability() {
    // Test without database
    let provider_no_db = SecretsProvider::new(None, None, None, "test-kms".to_string(), None, None);
    assert!(
        !provider_no_db.is_available(),
        "Should not be available without database"
    );

    // Try to set up with database
    let setup_result = setup_test_database().await;

    if setup_result.is_err() {
        info!("Database not available, skipping availability test with DB");
        return;
    }

    let (db, user_id, policy_client, _policy_data_address) = setup_result.unwrap();

    // Initialize global database
    let db_config = DatabaseConfig {
        url: std::env::var("DATABASE_URL")
            .unwrap_or_else(|_| "postgres://newton:newton@localhost:5432/newton_gateway".to_string()),
        max_connections: 5,
        min_connections: 1,
        connect_timeout: std::time::Duration::from_secs(10),
    };

    let _ = initialize_database(db_config).await;

    let provider_with_db = SecretsProvider::new(None, Some(policy_client), None, "test-kms".to_string(), None, None);

    assert!(provider_with_db.is_available(), "Should be available with database set");

    cleanup_test_database(&db, user_id).await.expect("Failed to cleanup");

    info!("SecretsProvider availability test completed");
}

/// Tests that multiple policy clients can have different encrypted values
/// for the same secret key. This validates:
/// 1. Secret keys are scoped per policy client
/// 2. Each client gets their own unique encrypted value
/// 3. Policy client isolation is properly maintained in the database
///
/// This is essential for multi-tenancy where different clients need different
/// API keys or credentials for the same service.
#[tokio::test]
async fn test_multiple_secrets_different_policy_clients() {
    let setup_result = setup_test_database().await;

    if setup_result.is_err() {
        info!("Database not available, skipping multiple policy clients test");
        return;
    }

    let (db, user_id, policy_client_1, policy_data_address) = setup_result.unwrap();

    // Create a second policy client with different secret
    let policy_client_2 = Address::from([0x43u8; 20]);

    // Claim ownership for the second policy client (required by FK).
    let owner_repo = PolicyClientOwnerRepository::new(db.clone());
    owner_repo
        .claim_if_unowned(user_id, policy_client_2)
        .await
        .expect("claim policy_client_2 ownership");

    sqlx::query(
        r#"
        INSERT INTO policy_client_secret (policy_client_address, policy_data_address, secrets)
        VALUES ($1, $2, $3)
        "#,
    )
    .bind(policy_client_2.as_slice())
    .bind(policy_data_address.as_slice())
    .bind(b"different_ciphertext".to_vec())
    .execute(db.pool())
    .await
    .expect("Failed to insert second secret");

    let repo = WasmSecretRepository::new(db.clone());

    // Get secrets blob for policy client 1
    let result_1 = repo
        .get_by_policy_data_and_policy_client(policy_data_address, policy_client_1)
        .await
        .expect("Query failed");
    assert!(result_1.is_some());

    // Get secrets blob for policy client 2
    let result_2 = repo
        .get_by_policy_data_and_policy_client(policy_data_address, policy_client_2)
        .await
        .expect("Query failed");
    assert!(result_2.is_some());

    cleanup_test_database(&db, user_id).await.expect("Failed to cleanup");

    info!("Multiple policy clients with different secrets test completed");
}

/// Tests that WasmExecutorConfig properly stores and provides access to the
/// KMS key ID configuration required for decrypting secrets. This validates
/// that the executor configuration includes all necessary information for the
/// secrets decryption infrastructure.
#[test]
fn test_wasm_executor_config_includes_secrets_config() {
    use newton_data_provider::WasmExecutorConfig;

    let config = WasmExecutorConfig {
        fuel: 100_000_000,
        max_wasm_stack: 64 * 1024 * 1024,
        max_http_calls: 50,
        max_http_request_size: 1024 * 1024,
        kms_rsa_key_id: "arn:aws:kms:us-west-2:123456789:key/test-key-id".to_string(),
    };

    assert!(!config.kms_rsa_key_id.is_empty(), "KMS key ID should be configured");
    assert!(
        config.kms_rsa_key_id.starts_with("arn:aws:kms"),
        "Should look like AWS KMS ARN"
    );

    info!("WasmExecutorConfig secrets configuration test completed");
}

fn load_test_secrets_schema_json() -> serde_json::Value {
    serde_json::from_str(include_str!("secrets_schema.json")).expect("invalid test secrets schema json")
}

async fn insert_plaintext_base64_secret(
    db: &DatabaseManager,
    user_id: Uuid,
    policy_client_address: Address,
    policy_data_address: Address,
    plaintext_value: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    let secrets_blob = STANDARD.encode(plaintext_value.as_bytes()).into_bytes();

    let owner_repo = PolicyClientOwnerRepository::new(db.clone());
    owner_repo
        .claim_if_unowned(user_id, policy_client_address)
        .await
        .map_err(|e| format!("failed to claim policy_client ownership: {e}"))?;

    sqlx::query(
        r#"
        INSERT INTO policy_client_secret (policy_client_address, policy_data_address, secrets)
        VALUES ($1, $2, $3)
        ON CONFLICT (policy_client_address, policy_data_address) DO UPDATE
        SET secrets = EXCLUDED.secrets
        "#,
    )
    .bind(policy_client_address.as_slice())
    .bind(policy_data_address.as_slice())
    .bind(secrets_blob)
    .execute(db.pool())
    .await?;

    Ok(())
}

async fn build_secrets_json_from_db(
    repo: &WasmSecretRepository,
    policy_client_address: Address,
    policy_data_address: Address,
) -> serde_json::Value {
    // In the real flow, the DB record contains an encrypted blob which decrypts to a JSON string:
    // { "COIN_GECKO_API": "...", "PERSONA_API": "..." }
    let rec = repo
        .get_by_policy_data_and_policy_client(policy_data_address, policy_client_address)
        .await
        .expect("db query failed")
        .expect("expected secret record to exist");

    let plaintext_bytes = STANDARD
        .decode(&rec.secrets_blob)
        .expect("secrets should be base64 string bytes for test");

    let plaintext_str = String::from_utf8(plaintext_bytes).expect("plaintext should be utf-8 for test");

    serde_json::from_str(&plaintext_str).expect("secret plaintext should be JSON")
}

#[tokio::test]
async fn test_secrets_schema_validation_multiple() {
    let setup_result = setup_test_database().await;
    if setup_result.is_err() {
        info!(
            "Database not available, skipping secrets schema validation DB test: {:?}",
            setup_result.err()
        );
        return;
    }

    let (db, user_id, policy_client_address, policy_data_address) = setup_result.unwrap();

    insert_plaintext_base64_secret(
        &db,
        user_id,
        policy_client_address,
        policy_data_address,
        r#"{"COIN_GECKO_API":"cg_key_123","PERSONA_API":"persona_key_456"}"#,
    )
    .await
    .unwrap();

    let repo = WasmSecretRepository::new(db.clone());
    let secrets = build_secrets_json_from_db(&repo, policy_client_address, policy_data_address).await;

    let schema = load_test_secrets_schema_json();
    newton_prover_core::rego::validate_secrets_schema(schema, secrets).unwrap();

    cleanup_test_database(&db, user_id).await.unwrap();
}

#[tokio::test]
async fn test_policy_multiple_policy_data_db_secrets_superset() {
    let setup_result = setup_test_database().await;
    if setup_result.is_err() {
        info!("Database not available, skipping multi-policy-data superset secrets test");
        return;
    }

    let (db, user_id, policy_client_address, policy_data_address) = setup_result.unwrap();

    // DB stores a merged superset for the policy client:
    // e.g. PolicyData1 reads COIN_GECKO_API, PolicyData2 reads PERSONA_API.
    insert_plaintext_base64_secret(
        &db,
        user_id,
        policy_client_address,
        policy_data_address,
        r#"{"COIN_GECKO_API":"cg_key_123","PERSONA_API":"persona_key_456","EXTRA":"ignored"}"#,
    )
    .await
    .unwrap();

    let repo = WasmSecretRepository::new(db.clone());
    let secrets = build_secrets_json_from_db(&repo, policy_client_address, policy_data_address).await;

    // Per-PolicyData schema #1 (could have required/additionalProperties, but we validate per-key only).
    let schema_policy_data_1 = serde_json::json!({
        "type": "object",
        "properties": {
            "COIN_GECKO_API": { "type": "string", "minLength": 1 }
        },
        "required": ["COIN_GECKO_API"],
        "additionalProperties": false
    });

    // Per-PolicyData schema #2
    let schema_policy_data_2 = serde_json::json!({
        "type": "object",
        "properties": {
            "PERSONA_API": { "type": "string", "minLength": 1 }
        },
        "required": ["PERSONA_API"],
        "additionalProperties": false
    });

    // Validate the full secrets object against each per-PolicyData schema (strict).
    newton_prover_core::rego::validate_secrets_schema(schema_policy_data_1, secrets.clone()).unwrap();
    newton_prover_core::rego::validate_secrets_schema(schema_policy_data_2, secrets.clone()).unwrap();

    // Negative check: schema constraint fails.
    let schema_too_strict = serde_json::json!({
        "type": "object",
        "properties": {
            "COIN_GECKO_API": { "type": "string", "minLength": 999 }
        },
        "required": ["COIN_GECKO_API"]
    });
    assert!(newton_prover_core::rego::validate_secrets_schema(schema_too_strict, secrets.clone()).is_err());

    cleanup_test_database(&db, user_id).await.unwrap();
}

#[tokio::test]
async fn test_secrets_schema_validation_from_database_rejects_missing_required_key() {
    let setup_result = setup_test_database().await;
    if setup_result.is_err() {
        info!("Database not available, skipping secrets schema missing-key DB test");
        return;
    }

    let (db, user_id, policy_client_address, policy_data_address) = setup_result.unwrap();

    insert_plaintext_base64_secret(
        &db,
        user_id,
        policy_client_address,
        policy_data_address,
        r#"{"COIN_GECKO_API":"cg_key_123"}"#,
    )
    .await
    .unwrap();

    let repo = WasmSecretRepository::new(db.clone());
    let secrets = build_secrets_json_from_db(&repo, policy_client_address, policy_data_address).await;

    let schema = load_test_secrets_schema_json();
    assert!(newton_prover_core::rego::validate_secrets_schema(schema, secrets).is_err());

    cleanup_test_database(&db, user_id).await.unwrap();
}

#[tokio::test]
async fn test_secrets_schema_validation_from_database_rejects_empty_value() {
    let setup_result = setup_test_database().await;
    if setup_result.is_err() {
        info!("Database not available, skipping secrets schema empty-value DB test");
        return;
    }

    let (db, user_id, policy_client_address, policy_data_address) = setup_result.unwrap();

    insert_plaintext_base64_secret(
        &db,
        user_id,
        policy_client_address,
        policy_data_address,
        r#"{"COIN_GECKO_API":"","PERSONA_API":"persona_key_456"}"#,
    )
    .await
    .unwrap();

    let repo = WasmSecretRepository::new(db.clone());
    let secrets = build_secrets_json_from_db(&repo, policy_client_address, policy_data_address).await;

    let schema = load_test_secrets_schema_json();
    assert!(newton_prover_core::rego::validate_secrets_schema(schema, secrets).is_err());

    cleanup_test_database(&db, user_id).await.unwrap();
}

#[tokio::test]
async fn test_secrets_schema_validation_from_database_rejects_additional_keys() {
    let setup_result = setup_test_database().await;
    if setup_result.is_err() {
        info!("Database not available, skipping secrets schema additional-keys DB test");
        return;
    }

    let (db, user_id, policy_client_address, policy_data_address) = setup_result.unwrap();

    insert_plaintext_base64_secret(
        &db,
        user_id,
        policy_client_address,
        policy_data_address,
        r#"{"COIN_GECKO_API":"cg_key_123","PERSONA_API":"persona_key_456","EXTRA":"nope"}"#,
    )
    .await
    .unwrap();

    let repo = WasmSecretRepository::new(db.clone());
    let secrets = build_secrets_json_from_db(&repo, policy_client_address, policy_data_address).await;

    let schema = load_test_secrets_schema_json();
    assert!(newton_prover_core::rego::validate_secrets_schema(schema, secrets).is_err());

    cleanup_test_database(&db, user_id).await.unwrap();
}
