use crate::{wasm::secrets::SecretsProvider, DataProviderConfig, DataSourceError};
use alloy::primitives::Address;
use cid::{multibase::Base, multihash::Multihash};
use newton_prover_core::{
    config::{ipfs::IpfsConfig, NewtonAvsConfig},
    database::try_get_database,
};
use ring::digest::{Context, SHA256};
use std::{
    net::IpAddr,
    sync::atomic::{AtomicU32, Ordering},
};
use url::Url;
use wasmtime::{
    component::{bindgen, Component, HasSelf, Linker, ResourceTable},
    Config, Engine,
};
use wasmtime_wasi::p2::{IoView, WasiCtx, WasiCtxBuilder, WasiView};

bindgen!({
    path: "wit",
    async: true,
});

use newton::provider::{
    http::{HttpRequest, HttpResponse},
    secrets::{Host as SecretsHost, SecretResponse},
};

impl SecretsHost for SecretsProvider {
    async fn get(&mut self) -> Result<SecretResponse, String> {
        let value = self.get_secrets_json().await?;
        Ok(SecretResponse { value })
    }
}

pub struct HttpProvider {
    call_count: AtomicU32,
    max_calls: u32,
    max_request_size: u32,
}

impl HttpProvider {
    fn new(max_calls: u32, max_request_size: u32) -> Self {
        Self {
            call_count: AtomicU32::default(),
            max_calls,
            max_request_size,
        }
    }
}

impl newton::provider::http::Host for HttpProvider {
    async fn fetch(&mut self, request: HttpRequest) -> Result<HttpResponse, String> {
        // limit number of http calls wasm is allowed to make
        let current_count = self.call_count.fetch_add(1, Ordering::Relaxed);
        if current_count >= self.max_calls {
            return Err(format!("exceeded maximum HTTP calls limit of {}", self.max_calls));
        }

        // limit body + headers request size
        let headers_size: usize = request
            .headers
            .iter()
            .map(|(k, v)| k.len() + v.len() + 4) // +4 for ": " and "\r\n"
            .sum();
        let body_size = request.body.as_ref().map(|b| b.len()).unwrap_or(0);
        let total_size = headers_size + body_size;

        if total_size > self.max_request_size as usize {
            return Err(format!(
                "request size {} exceeds maximum of {} bytes",
                total_size, self.max_request_size
            ));
        }

        let url = Url::parse(&request.url).map_err(|e| format!("invalid URL: {e}"))?;

        // deny access to private IP blocks
        if let Some(host_str) = url.host_str() {
            if let Ok(ip) = host_str.parse::<IpAddr>() {
                match ip {
                    IpAddr::V4(ipv4) => {
                        // deny private v4 blocks
                        if ipv4.is_private() || ipv4.is_loopback() || ipv4.is_link_local() ||
                           ipv4.is_broadcast() || ipv4.is_multicast() ||
                           // and aws
                           (ipv4.octets()[0] == 169 && ipv4.octets()[1] == 254)
                        {
                            return Err("access to private/reserved IP ranges is blocked".to_string());
                        }
                    }
                    IpAddr::V6(ipv6) => {
                        // deny private IPv6 blocks
                        if ipv6.is_loopback() || ipv6.is_multicast() ||
                           // private/local ranges
                           ipv6.segments()[0] & 0xfe00 == 0xfc00 || // fc00::/7 local
                           ipv6.segments()[0] & 0xffc0 == 0xfe80
                        {
                            // fe80::/10 linklocal
                            return Err("access to private/reserved IP ranges is blocked".to_string());
                        }
                    }
                }
            }
        }

        // WIT -> reqwest request
        let method = match request.method.as_str() {
            "GET" => reqwest::Method::GET,
            "POST" => reqwest::Method::POST,
            "PUT" => reqwest::Method::PUT,
            "DELETE" => reqwest::Method::DELETE,
            _ => return Err(format!("unsupported HTTP method: {}", request.method)),
        };

        let client = reqwest::Client::new();
        let mut req_builder = client.request(method, &request.url);

        // add headers
        for (key, value) in request.headers {
            req_builder = req_builder.header(&key, &value);
        }

        // add body if present
        if let Some(body) = request.body {
            req_builder = req_builder.body(body);
        }

        // execute request
        match req_builder.send().await {
            Ok(response) => {
                let status = response.status().as_u16();
                let headers = response
                    .headers()
                    .iter()
                    .map(|(k, v)| (k.to_string(), v.to_str().unwrap_or_default().to_string()))
                    .collect();

                match response.bytes().await {
                    Ok(body) => Ok(HttpResponse {
                        status,
                        headers,
                        body: body.to_vec(),
                    }),
                    Err(e) => Err(format!("failed to read response body: {e}")),
                }
            }
            Err(e) => Err(format!("HTTP request failed: {e}")),
        }
    }
}

/// Configuration for WASM execution
#[derive(Debug, Clone)]
pub struct WasmExecutorConfig {
    /// Instruction limit (fuel)
    pub fuel: u64,
    /// Max WASM stack size in bytes
    pub max_wasm_stack: u64,
    /// Max HTTP calls per execution
    pub max_http_calls: u32,
    /// Max HTTP request size (body + headers) in bytes
    pub max_http_request_size: u32,
    /// KMS RSA key ID for decrypting secrets
    pub kms_rsa_key_id: String,
}

impl WasmExecutorConfig {
    /// Create config from DataProviderConfig
    pub fn from_data_provider_config(config: &NewtonAvsConfig<DataProviderConfig>) -> Self {
        Self {
            fuel: config.service.fuel,
            max_wasm_stack: config.service.max_wasm_stack,
            max_http_calls: config.service.max_http_calls,
            max_http_request_size: config.service.max_http_request_size,
            kms_rsa_key_id: config.service.kms_rsa_key_id.clone(),
        }
    }
}

#[derive(Debug)]
pub struct WasmExecutor {
    engine: Engine,
    config: WasmExecutorConfig,
}

impl WasmExecutor {
    /// Create executor from DataProviderConfig (backwards compatibility)
    pub fn from_config(config: &NewtonAvsConfig<DataProviderConfig>) -> Result<Self, DataSourceError> {
        Self::new(WasmExecutorConfig::from_data_provider_config(config))
    }

    /// Create executor with custom configuration
    pub fn new(config: WasmExecutorConfig) -> Result<Self, DataSourceError> {
        let mut wasm_config = Config::new();
        wasm_config.wasm_component_model(true);
        wasm_config.async_support(true);

        // prevent resource exhaustion
        wasm_config.consume_fuel(config.fuel > 0);
        wasm_config.max_wasm_stack(config.max_wasm_stack as usize);

        // Configures the size of the stacks used for asynchronous execution.
        //
        // This setting configures the size of the stacks that are allocated for asynchronous
        // execution. The value cannot be less than max_wasm_stack. The amount of stack space
        // guaranteed for host functions is async_stack_size - max_wasm_stack, so take care not to
        // set these two values close to one another; doing so may cause host functions to overflow
        // the stack and abort the process.
        wasm_config.async_stack_size(config.max_wasm_stack.saturating_mul(2) as usize);

        let engine = Engine::new(&wasm_config)
            .map_err(|e| DataSourceError::ParseError(format!("failed to create wasm engine: {e}")))?;

        Ok(Self { engine, config })
    }

    pub fn hash_wasm_binary(wasm_bytes: &[u8]) -> String {
        let mut context = Context::new(&SHA256);
        context.update(wasm_bytes);
        let digest = context.finish();

        let m_hash: Multihash<64> = Multihash::wrap(0x12, digest.as_ref()).expect("failed to wrap hash");
        cid::multibase::encode(Base::Base16Lower, m_hash.to_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<IpfsConfig>,
    ) -> Result<Result<String, String>, DataSourceError> {
        let component = Component::new(&self.engine, wasm_bytes)
            .map_err(|e| DataSourceError::ParseError(format!("failed to compile wasm component: {e}")))?;
        struct MyCtx {
            table: ResourceTable,
            wasi: WasiCtx,
            http: HttpProvider,
            secrets: SecretsProvider,
        }

        impl IoView for MyCtx {
            fn table(&mut self) -> &mut ResourceTable {
                &mut self.table
            }
        }

        impl WasiView for MyCtx {
            fn ctx(&mut self) -> &mut WasiCtx {
                &mut self.wasi
            }
        }

        // wasi context
        let db_pool = try_get_database().map(|db| db.pool().clone());
        let ctx = MyCtx {
            table: ResourceTable::new(),
            secrets: SecretsProvider::new(
                db_pool,
                policy_client_address,
                policy_data_address,
                self.config.kms_rsa_key_id.clone(),
                secrets_schema_cid,
                ipfs_config,
            ),
            http: HttpProvider::new(self.config.max_http_calls, self.config.max_http_request_size),
            wasi: WasiCtxBuilder::new()
                .args(&["plugin", input]) // pass input as command line argument
                .inherit_stdin()
                .inherit_stdout()
                .inherit_stderr()
                // NO network access - use HttpProvider only
                // NO env access at all
                .build(),
        };

        let mut store = wasmtime::Store::new(&self.engine, ctx);

        // attempting to `set_fuel` when fuel has not been enabled results in an err so don't do
        // that
        if self.config.fuel > 0 {
            store
                .set_fuel(self.config.fuel)
                .map_err(|e| DataSourceError::ParseError(format!("failed to set fuel: {e}")))?;
        }

        let mut linker = Linker::new(&self.engine);
        linker.allow_shadowing(true);

        // add wasi functions to linker
        wasmtime_wasi::p2::add_to_linker_async(&mut linker)
            .map_err(|e| DataSourceError::ParseError(format!("failed to add wasi to linker: {e}")))?;

        // add the newton provider interface to linker
        newton::provider::http::add_to_linker::<MyCtx, HasSelf<HttpProvider>>(&mut linker, |ctx: &mut MyCtx| {
            &mut ctx.http
        })
        .map_err(|e| DataSourceError::ParseError(format!("failed to add newton provider to linker: {e}")))?;

        // add the secrets provider interface to linker (provider handles availability internally)
        newton::provider::secrets::add_to_linker::<MyCtx, HasSelf<SecretsProvider>>(&mut linker, |ctx: &mut MyCtx| {
            &mut ctx.secrets
        })
        .map_err(|e| DataSourceError::ParseError(format!("failed to add secrets provider to linker: {e}")))?;

        // instantiate the newton provider world
        let newton_provider = NewtonProvider::instantiate_async(&mut store, &component, &linker)
            .await
            .map_err(|e| DataSourceError::ParseError(format!("failed to instantiate wasm component: {e}")))?;

        // execute the run function with JSON input
        let result = newton_provider
            .call_run(&mut store, input)
            .await
            .map_err(|e| DataSourceError::ParseError(format!("wasm execution failed: {e}")))?;

        Ok(result)
    }
}
