#[macro_use]
extern crate newton_prover_core;

use alloy::{
    primitives::{keccak256, Address, Bytes, B256},
    providers::Provider,
    signers::{local::PrivateKeySigner, Signer},
    sol_types::SolValue,
};
use eigensdk::common::get_provider;
use newton_prover_core::{
    config::{ipfs::IpfsConfig, NewtonAvsConfig},
    newton_policy_data::{INewtonPolicyData::AttestationInfo, NewtonPolicyData},
    newton_prover_task_manager::NewtonMessage,
};
use serde::{Deserialize, Serialize};
use serde_json::Value as JsonValue;
use std::collections::HashMap;
use thiserror::Error;

pub mod config;
pub mod data_cache;
pub mod data_sources;
pub mod policy_data;
pub mod wasm;
pub mod wasm_cache;
pub mod wasm_provider;

pub use config::DataProviderConfig;
pub use data_cache::{PolicyTaskDataCache, RedisCacheError};
pub use data_sources::DataSource;
pub use policy_data::PolicyDataWasm;
use tracing::{info, warn};
pub use wasm::{WasmExecutor, WasmExecutorConfig};
pub use wasm_cache::WasmCache;
pub use wasm_provider::WasmProvider;

#[derive(Debug, Error)]
pub enum DataSourceError {
    #[error("Unsupported provider: {0}")]
    UnsupportedProvider(String),
    #[error("Network error: {0}")]
    NetworkError(String),
    #[error("Parse error: {0}")]
    ParseError(String),
    #[error("Authentication error: {0}")]
    AuthError(String),
    #[error("Rate limit exceeded")]
    RateLimitExceeded,
    #[error("Signer error: {0}")]
    SignerError(String),
    #[error("Redis cache error: {0}")]
    RedisCacheError(#[from] RedisCacheError),
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct PolicyDataInfo {
    policy_data_address: Address,
    wasm_cid: String,
    secrets_schema_cid: String,
    attestation_info: AttestationInfo,
}

#[derive(Debug)]
pub struct DataProviderBuilder {
    config: NewtonAvsConfig<DataProviderConfig>,
    data_cache_url: Option<String>,
}

impl DataProviderBuilder {
    pub fn new(config: NewtonAvsConfig<DataProviderConfig>) -> Self {
        Self {
            config,
            data_cache_url: None,
        }
    }

    pub fn with_cache(&mut self, data_cache_url: &str) -> &mut Self {
        self.data_cache_url = Some(data_cache_url.to_string());
        self
    }

    pub async fn build(&self) -> Result<DataProvider, DataSourceError> {
        let mut data_provider = DataProvider {
            rpc_url: self.config.rpc.http.clone(),
            data_source: DataSource::from_config(&self.config).unwrap(),
            data_cache: self
                .data_cache_url
                .as_ref()
                .map(|url| PolicyTaskDataCache::new(url).unwrap()),
            chain_id: self.config.chain_id,
            ipfs_config: self.config.ipfs.clone(),
        };
        let _ = data_provider.initialize().await.inspect_err(|e| {
            warn!("Failed to initialize data provider: {}", e);
        });
        Ok(data_provider)
    }
}

#[derive(Debug)]
pub struct DataProvider {
    rpc_url: String,
    data_source: DataSource,
    data_cache: Option<PolicyTaskDataCache>,
    chain_id: u64,
    ipfs_config: IpfsConfig,
}

impl DataProvider {
    pub async fn run_plugin(
        &mut self,
        provider: WasmProvider,
        policy_client_address: Option<Address>,
        policy_data_address: Option<Address>,
        secrets_schema_cid: Option<String>,
    ) -> Result<PolicyDataWasm, DataSourceError> {
        self.data_source
            .fetch_data(
                provider,
                policy_client_address,
                policy_data_address,
                secrets_schema_cid,
                Some(self.ipfs_config.clone()),
            )
            .await
    }

    pub async fn initialize(&mut self) -> Result<(), DataSourceError> {
        self.data_source.initialize().await?;

        info!("Initializing data provider (with cache: {})", self.data_cache.is_some());
        if let Some(ref mut data_cache) = self.data_cache {
            match data_cache.initialize().await {
                Ok(()) => info!("Redis cache initialized successfully"),
                Err(e) => {
                    warn!(error = %e, "Failed to initialize Redis cache, continuing without cache");
                    // Set cache to None so it won't be used
                    self.data_cache = None;
                }
            }
        }

        Ok(())
    }

    /// Fetches and merges all secrets schemas (by `secretsSchemaCid`) for the policy set on a policy client.
    ///
    /// Merge rules:
    /// - `properties`: union keys (first definition wins for duplicates)
    /// - `required`: union
    /// - `type`: `"object"`
    /// - `additionalProperties`: `false` if any schema sets it to `false`; otherwise preserve `true`/unset
    pub async fn get_merged_secrets_schema_for_client(
        &self,
        policy_client_address: Address,
    ) -> Result<JsonValue, DataSourceError> {
        use newton_prover_core::{newton_policy::NewtonPolicy, newton_policy_client::NewtonPolicyClient};

        let provider = get_provider(&self.rpc_url);

        let policy_client = NewtonPolicyClient::new(policy_client_address, &provider);
        let policy_address = policy_client
            .getPolicyAddress()
            .call()
            .await
            .map_err(|e| DataSourceError::NetworkError(format!("Failed to get policy address: {e}")))?;

        let policy = NewtonPolicy::new(policy_address, &provider);
        let policy_data_addresses = policy
            .getPolicyData()
            .call()
            .await
            .map_err(|e| DataSourceError::NetworkError(format!("Failed to get policy data addresses: {e}")))?;

        let mut schema_cids: Vec<String> = Vec::new();
        for policy_data_address in policy_data_addresses {
            let policy_data = NewtonPolicyData::new(policy_data_address, &provider);
            let cid = policy_data.getSecretsSchemaCid().call().await.map_err(|e| {
                DataSourceError::NetworkError(format!(
                    "Failed to get secrets schema cid for policy data {policy_data_address}: {e}"
                ))
            })?;

            if !cid.is_empty() {
                schema_cids.push(cid);
            }
        }

        if schema_cids.is_empty() {
            return Err(DataSourceError::ParseError(
                "No secretsSchemaCid configured on any PolicyData for this policy".to_string(),
            ));
        }

        // Optional per-call CID de-dupe to avoid double fetches (first-seen wins globally anyway).
        let mut seen_cids: HashMap<String, ()> = HashMap::new();
        let mut schema_docs: Vec<(String, JsonValue)> = Vec::new();
        for cid in schema_cids {
            if seen_cids.insert(cid.clone(), ()).is_some() {
                continue;
            }
            let schema_text = newton_prover_core::common::fetch_from_ipfs(&cid, Some(&self.ipfs_config))
                .await
                .map_err(|e| {
                    DataSourceError::NetworkError(format!("Failed to fetch secrets schema from IPFS (cid={cid}): {e}"))
                })?
                .text()
                .await
                .map_err(|e| {
                    DataSourceError::NetworkError(format!("Failed to read secrets schema text (cid={cid}): {e}"))
                })?;

            let schema_json: JsonValue = serde_json::from_str(&schema_text)
                .map_err(|e| DataSourceError::ParseError(format!("Invalid secrets schema JSON (cid={cid}): {e}")))?;
            schema_docs.push((cid, schema_json));
        }
        newton_prover_core::common::task::merge_secrets_schemas(schema_docs)
            .map_err(|e| DataSourceError::ParseError(e.to_string()))
    }

    /// get policy information from contracts
    async fn get_policy_for_client(
        &self,
        policy_client_address: Address,
    ) -> Result<(Address, B256, String, u32, Option<PolicyDataInfo>), DataSourceError> {
        use newton_prover_core::{newton_policy::NewtonPolicy, newton_policy_client::NewtonPolicyClient};

        let provider = get_provider(&self.rpc_url);

        info!(
            "[get_policy_for_client] policy_client_address: {}",
            policy_client_address
        );

        // get policy address from policy client
        let policy_client = NewtonPolicyClient::new(policy_client_address, &provider);
        let policy_address = policy_client
            .getPolicyAddress()
            .call()
            .await
            .map_err(|e| DataSourceError::NetworkError(e.to_string()))?;

        info!("[get_policy_for_client] policy_address: {}", policy_address);

        // get policy data address from policy
        let policy = NewtonPolicy::new(policy_address, &provider);
        info!("[get_policy_for_client] policy: {}", policy_address);
        let policy_data_address = policy
            .getPolicyData()
            .call()
            .await
            .map_err(|e| DataSourceError::NetworkError(format!("Failed to get policy data: {}", e)))?;
        info!(
            "[get_policy_for_client] policy_data_address: {}",
            policy_data_address
                .iter()
                .map(|a| a.to_string())
                .collect::<Vec<_>>()
                .join(", ")
        );
        let policy_id = policy_client
            .getPolicyId()
            .call()
            .await
            .map_err(|e| DataSourceError::NetworkError(format!("Failed to get policy id: {}", e)))?;
        info!("[get_policy_for_client] policy_id: {}", policy_id);

        let policy_url = policy
            .getPolicyCid()
            .call()
            .await
            .map_err(|e| DataSourceError::NetworkError(format!("Failed to get policy cid: {}", e)))?;

        let current_block = provider
            .get_block_number()
            .await
            .map_err(|e| DataSourceError::NetworkError(e.to_string()))? as u32;
        let expire_after = policy
            .getPolicyConfig(policy_id)
            .call()
            .await
            .map_err(|e| DataSourceError::NetworkError(format!("Failed to get policy config: {}", e)))?
            .expireAfter;
        let expire_block = current_block + expire_after;

        // TODO: @weslie handle multiple policy data addresses
        // get url + args from policy data contract
        let policy_data_info = if !policy_data_address.is_empty() {
            info!(
                "[get_policy_for_client] policy_data_address: {}",
                policy_data_address[0]
            );
            Some(self.get_policy_data_for_policy(policy_data_address[0]).await?)
        } else {
            None
        };

        Ok((policy_address, policy_id, policy_url, expire_block, policy_data_info))
    }

    async fn get_policy_data_for_policy(
        &self,
        policy_data_address: Address,
    ) -> Result<PolicyDataInfo, DataSourceError> {
        let provider = get_provider(&self.rpc_url);
        let policy_data = NewtonPolicyData::new(policy_data_address, &provider);
        let wasm_cid = policy_data
            .getWasmCid()
            .call()
            .await
            .map_err(|e| DataSourceError::NetworkError(format!("Failed to get policy data location: {}", e)))?;

        // failure to fetch `secretsSchemaCid` should not fail policy data fetching.
        // Treat as empty CID (no schema) and continue.
        let secrets_schema_cid = match policy_data.getSecretsSchemaCid().call().await {
            Ok(cid) => cid,
            Err(e) => {
                warn!(
                    "Failed to get secrets schema cid for policy_data_address {}: {}",
                    policy_data_address, e
                );
                String::new()
            }
        };

        let attestation_info = policy_data
            .getAttestationInfo()
            .call()
            .await
            .map_err(|e| DataSourceError::NetworkError(format!("Failed to get attestation info: {}", e)))?;

        Ok(PolicyDataInfo {
            policy_data_address,
            wasm_cid,
            secrets_schema_cid,
            attestation_info,
        })
    }

    /// fetch policy data from urls and process through wasm
    async fn fetch_policy(&mut self, policy_url: &str) -> Result<String, DataSourceError> {
        let policy_text = newton_prover_core::common::fetch_from_ipfs(policy_url, Some(&self.ipfs_config))
            .await
            .map_err(|e| DataSourceError::NetworkError(e.to_string()))?
            .text()
            .await
            .map_err(|e| DataSourceError::NetworkError(e.to_string()))?;
        Ok(policy_text)
    }

    /// fetch policy data from urls and process through wasm
    async fn execute_policy_data_wasm(
        &mut self,
        wasm_url: &str,
        wasm_args: Option<Bytes>,
        policy_client_address: Address,
        policy_data_address: Address,
        secrets_schema_cid: Option<String>,
    ) -> Result<PolicyDataWasm, DataSourceError> {
        let wasm_provider = WasmProvider::with_args(wasm_url, wasm_args.unwrap_or_default());

        let policy_data_result = self
            .run_plugin(
                wasm_provider,
                Some(policy_client_address),
                Some(policy_data_address),
                secrets_schema_cid,
            )
            .await?;

        Ok(policy_data_result)
    }

    /// will fetch data needed to evaluate the given policy
    pub async fn get_policy_task_data_for_client(
        &mut self,
        signer: &PrivateKeySigner,
        policy_client_address: Address,
        wasm_args: Option<Bytes>,
    ) -> Result<NewtonMessage::PolicyTaskData, DataSourceError> {
        let chain_id = self.chain_id;
        let (policy_address, policy_id, policy_url, expire_block, policy_data_info) =
            self.get_policy_for_client(policy_client_address).await?;

        let policy_data = if let Some(policy_data_info) = &policy_data_info {
            // Check cache first if Redis is available
            if let Some(ref data_cache) = self.data_cache {
                if data_cache.is_available().await {
                    match data_cache
                        .get(
                            policy_client_address,
                            policy_data_info.policy_data_address,
                            chain_id,
                            wasm_args.clone(),
                        )
                        .await
                    {
                        Ok(cached_data) => {
                            info!(
                                "Returning cached policy task data for address {}",
                                policy_data_info.policy_data_address
                            );
                            return Ok(cached_data);
                        }
                        Err(RedisCacheError::CacheMiss) => {
                            info!(
                                "Cache miss for policy data address {}, fetching fresh data",
                                policy_data_info.policy_data_address
                            );
                        }
                        Err(e) => {
                            info!("Redis cache error: {}, falling back to fresh data fetch", e);
                        }
                    }
                }
            }

            let secrets_schema_cid = if !policy_data_info.secrets_schema_cid.is_empty() {
                Some(policy_data_info.secrets_schema_cid.clone())
            } else {
                None
            };

            let policy_data_wasm = self
                .execute_policy_data_wasm(
                    &policy_data_info.wasm_cid,
                    wasm_args.clone(),
                    policy_client_address,
                    policy_data_info.policy_data_address,
                    secrets_schema_cid,
                )
                .await?;
            let policy_data_bytes = serde_json::to_string(&policy_data_wasm.data).unwrap().into_bytes();

            let wasm_args_bytes = wasm_args.clone().unwrap_or_default();

            let message_hash = keccak256(
                (
                    wasm_args_bytes.as_ref(),
                    policy_data_bytes.as_slice(),
                    policy_data_info.policy_data_address,
                    expire_block,
                    policy_data_info.wasm_cid.clone(),
                    policy_data_info.secrets_schema_cid.clone(),
                )
                    .abi_encode_packed(),
            );

            info!("message_hash: {}", hex!(message_hash.as_slice()));

            let attestation = signer
                .sign_hash(&message_hash)
                .await
                .map_err(|e| DataSourceError::SignerError(e.to_string()))?;

            info!("attestation: {}", hex!(attestation.as_bytes()));

            Some(NewtonMessage::PolicyData {
                wasmArgs: wasm_args_bytes,
                data: Bytes::from(policy_data_bytes),
                attestation: Bytes::from(attestation.as_bytes()),
                policyDataAddress: policy_data_info.policy_data_address,
                expireBlock: expire_block,
            })
        } else {
            warn!("No policy data info found. Policy with no policy data found.");
            None
        };

        let policy = self.fetch_policy(policy_url.as_str()).await?;

        let policy_task_data = NewtonMessage::PolicyTaskData {
            policyId: policy_id,
            policyAddress: policy_address,
            policy: Bytes::from(policy.into_bytes()),
            policyData: match policy_data {
                Some(policy_data) => vec![policy_data],
                None => vec![],
            },
        };

        if let Some(policy_data_info) = policy_data_info {
            // Cache the result if Redis is available
            if let Some(ref data_cache) = self.data_cache {
                if data_cache.is_available().await {
                    // Calculate expire_after_blocks from the expire_block and current block
                    let provider = get_provider(&self.rpc_url);
                    let current_block = provider
                        .get_block_number()
                        .await
                        .map_err(|e| DataSourceError::NetworkError(e.to_string()))?
                        as u32;
                    let expire_after_blocks = expire_block.saturating_sub(current_block);

                    if let Err(e) = data_cache
                        .set(
                            policy_client_address,
                            policy_data_info.policy_data_address,
                            wasm_args,
                            policy_task_data.clone(),
                            expire_after_blocks,
                            self.chain_id,
                        )
                        .await
                    {
                        info!("Failed to cache policy task data: {}", e);
                    }
                }
            }
        }

        Ok(policy_task_data)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::wasm::secrets::SecretsProvider;
    use alloy::primitives::Address;

    /// Tests that SecretsProvider returns an appropriate error when attempting
    /// to retrieve secrets without a database connection initialized
    #[tokio::test]
    async fn test_secrets_provider_unavailable_without_database() {
        let policy_client_address = Address::from([1u8; 20]);
        let mut secrets_provider = SecretsProvider::new(
            None,
            Some(policy_client_address),
            None,
            "test-kms-key-id".to_string(),
            Some("bafk-test-schema-cid".to_string()),
            None,
        );

        let result = secrets_provider.get_secrets_json().await;

        assert!(result.is_err());
        assert_eq!(result.unwrap_err(), "Secrets are not available in this context");
    }

    /// Tests that SecretsProvider.is_available() correctly returns false
    /// when no database connection is configured
    #[tokio::test]
    async fn test_secrets_provider_is_not_available_without_database() {
        let policy_client_address = Address::from([2u8; 20]);
        let secrets_provider = SecretsProvider::new(
            None,
            Some(policy_client_address),
            None,
            "test-kms-key-id".to_string(),
            None,
            None,
        );

        assert!(!secrets_provider.is_available());
    }

    /// Tests that WasmExecutorConfig can be created with a KMS key ID
    /// for secrets decryption functionality
    #[tokio::test]
    async fn test_wasm_executor_config_direct_creation() {
        let config = WasmExecutorConfig {
            fuel: 100_000_000,
            max_wasm_stack: 1024 * 1024 * 64,
            max_http_calls: 50,
            max_http_request_size: 1024 * 1024,
            kms_rsa_key_id: "test-kms-key-123".to_string(),
        };

        assert_eq!(config.fuel, 100_000_000);
        assert_eq!(config.max_wasm_stack, 1024 * 1024 * 64);
        assert_eq!(config.max_http_calls, 50);
        assert_eq!(config.max_http_request_size, 1024 * 1024);
        assert_eq!(config.kms_rsa_key_id, "test-kms-key-123");
    }

    /// Tests that WasmExecutorConfig can be created with various default values
    /// including an empty KMS key ID
    #[test]
    fn test_wasm_executor_config_defaults() {
        let config = WasmExecutorConfig {
            fuel: 50_000_000,
            max_wasm_stack: 1024 * 1024 * 32,
            max_http_calls: 25,
            max_http_request_size: 512 * 1024,
            kms_rsa_key_id: String::new(),
        };

        assert_eq!(config.fuel, 50_000_000);
        assert_eq!(config.max_wasm_stack, 1024 * 1024 * 32);
        assert_eq!(config.max_http_calls, 25);
        assert_eq!(config.max_http_request_size, 512 * 1024);
        assert_eq!(config.kms_rsa_key_id, "");
    }

    /// Tests that SecretsProvider can be created with different policy client addresses,
    /// ensuring proper scoping for multi-tenant secret isolation
    #[tokio::test]
    async fn test_secrets_provider_policy_client_address_scoping() {
        let policy_client_1 = Address::from([1u8; 20]);
        let policy_client_2 = Address::from([2u8; 20]);

        let secrets_provider_1 = SecretsProvider::new(
            None,
            Some(policy_client_1),
            None,
            "test-kms-key".to_string(),
            None,
            None,
        );

        let secrets_provider_2 = SecretsProvider::new(
            None,
            Some(policy_client_2),
            None,
            "test-kms-key".to_string(),
            None,
            None,
        );

        assert!(!secrets_provider_1.is_available());
        assert!(!secrets_provider_2.is_available());
    }

    /// Tests that SecretsProvider can be created without a specific policy client address,
    /// allowing for global secrets that aren't scoped to a particular client
    #[tokio::test]
    async fn test_secrets_provider_without_policy_client_address() {
        let secrets_provider = SecretsProvider::new(None, None, None, "test-kms-key".to_string(), None, None);

        assert!(!secrets_provider.is_available());
    }

    /// Tests that DataProviderConfig allows setting and updating the KMS key ID
    /// used for decrypting WASM secrets
    #[tokio::test]
    async fn test_data_provider_config_kms_key_id_setters() {
        let mut config = DataProviderConfig::default();

        config.set_kms_rsa_key_id("updated-kms-key-456".to_string());

        assert_eq!(config.kms_rsa_key_id, "updated-kms-key-456");
    }

    #[test]
    fn merge_secrets_schemas_unions_required_and_properties_first_wins() {
        let schema_a = serde_json::json!({
            "type": "object",
            "properties": {
                "COIN_GECKO_API": { "type": "string", "minLength": 1 }
            },
            "required": ["COIN_GECKO_API"],
            "additionalProperties": false
        });

        // Duplicate key with different constraint should be ignored (first wins)
        let schema_b = serde_json::json!({
            "type": "object",
            "properties": {
                "COIN_GECKO_API": { "type": "string", "minLength": 999 },
                "PERSONA_API": { "type": "string", "minLength": 1 }
            },
            "required": ["PERSONA_API"]
        });

        let merged = newton_prover_core::common::task::merge_secrets_schemas(vec![
            ("cid_a".to_string(), schema_a),
            ("cid_b".to_string(), schema_b),
        ])
        .expect("merge");

        let props = merged
            .get("properties")
            .and_then(|v| v.as_object())
            .expect("properties object");

        assert!(props.contains_key("COIN_GECKO_API"));
        assert!(props.contains_key("PERSONA_API"));

        // First schema wins for duplicate property definitions
        assert_eq!(
            props.get("COIN_GECKO_API").unwrap().get("minLength").unwrap(),
            &serde_json::json!(1)
        );

        let required = merged
            .get("required")
            .and_then(|v| v.as_array())
            .expect("required array");
        assert!(required.contains(&serde_json::json!("COIN_GECKO_API")));
        assert!(required.contains(&serde_json::json!("PERSONA_API")));

        assert!(merged.get("additionalProperties").is_none());
    }
}
