use crate::{DataProviderConfig, DataSourceError, PolicyDataWasm, WasmCache, WasmExecutor, WasmProvider};
use alloy::primitives::Address;
use directories::ProjectDirs;
use newton_prover_core::config::NewtonAvsConfig;
use serde_json::Value;
use std::path::PathBuf;
use tracing::info;

/// Data source that executes Wasm plugins for data retrieval
#[derive(Debug)]
pub struct DataSource {
    executor: WasmExecutor,
    cache: WasmCache,
}

impl DataSource {
    pub fn from_config(config: &NewtonAvsConfig<DataProviderConfig>) -> Result<Self, DataSourceError> {
        let cache = WasmCache::new(
            ProjectDirs::from("", "", "newton-prover")
                .map(|dirs| dirs.cache_dir().join("wasm"))
                .unwrap_or_else(|| PathBuf::from("./cache/wasm")),
            config.service.max_cache_size,
            config.ipfs.clone(),
        );

        Ok(Self {
            executor: WasmExecutor::from_config(config)?,
            cache,
        })
    }

    pub async fn initialize(&mut self) -> Result<(), DataSourceError> {
        self.cache.initialize().await
    }

    pub fn hash_wasm_binary(wasm_bytes: &[u8]) -> String {
        WasmExecutor::hash_wasm_binary(wasm_bytes)
    }

    pub async fn execute_wasm_bytes(
        &self,
        wasm_bytes: &[u8],
        input: &str,
        policy_client_address: Option<Address>,
        policy_data_address: Option<Address>,
        secrets_schema_cid: Option<String>,
        ipfs_config: Option<newton_prover_core::config::ipfs::IpfsConfig>,
    ) -> Result<Result<String, String>, DataSourceError> {
        self.executor
            .execute_wasm_bytes(
                wasm_bytes,
                input,
                policy_client_address,
                policy_data_address,
                secrets_schema_cid,
                ipfs_config,
            )
            .await
    }
}

impl DataSource {
    pub async fn fetch_data(
        &mut self,
        provider: WasmProvider,
        policy_client_address: Option<Address>,
        policy_data_address: Option<Address>,
        secrets_schema_cid: Option<String>,
        ipfs_config: Option<newton_prover_core::config::ipfs::IpfsConfig>,
    ) -> Result<PolicyDataWasm, DataSourceError> {
        let input_str = provider.args_as_str();
        info!(
            url = provider.url,
            args = provider.args_as_str(),
            "executing data provider plugin"
        );

        // get wasm from cache or download
        let wasm_bytes = self.cache.get(&provider.url).await?;

        // execute wasm plugin
        let execution_result = self
            .execute_wasm_bytes(
                &wasm_bytes,
                &input_str,
                policy_client_address,
                policy_data_address,
                secrets_schema_cid,
                ipfs_config,
            )
            .await?;

        // handle the result from the wasm function
        let output_str = match execution_result {
            Ok(output) => output,
            Err(error_msg) => {
                return Err(DataSourceError::ParseError(format!("Data source error: {error_msg}")));
            }
        };

        // parse output JSON
        let data: Value = serde_json::from_str(&output_str)
            .map_err(|e| DataSourceError::ParseError(format!("failed to parse wasm output: {e}")))?;

        info!("wasm output: {}", output_str);

        Ok(PolicyDataWasm::new(provider, data))
    }
}
