use alloy::primitives::{Address, FixedBytes};
use clap::{Parser, Subcommand};
use eyre::{Context, Result};
use newton_prover_chainio::policy_data::PolicyDataController;
use newton_prover_core::{
    common::chain::{get_block_time_ms, get_chain_id},
    config::NewtonAvsConfig,
    newton_policy_data::INewtonPolicyData,
};
use serde_json::Value;
use std::path::PathBuf;
use tracing::{self, info};

use crate::{commands::utils, config::NewtonCliConfig};

/// Policy data commands
#[derive(Debug, Parser)]
#[command(name = "policy-data")]
pub struct PolicyDataCommand {
    #[command(subcommand)]
    pub subcommand: PolicyDataSubcommand,
}

#[derive(Debug, Subcommand)]
pub enum PolicyDataSubcommand {
    /// Deploy policy data
    Deploy(DeployCommand),
    /// Simulate WASM execution
    Simulate(SimulateCommand),
}

/// Simulate WASM execution command
#[derive(Debug, Parser)]
pub struct SimulateCommand {
    /// Path to the WASM component file
    #[arg(long)]
    wasm_file: PathBuf,

    /// Input JSON string
    #[arg(long)]
    input_json: String,
}

/// Deploy policy data command
#[derive(Debug, Parser)]
pub struct DeployCommand {
    #[arg(long, env = "PRIVATE_KEY")]
    private_key: Option<String>,

    #[arg(long, env = "RPC_URL")]
    rpc_url: Option<String>,

    #[arg(long, default_value = "policy-files/policy_cids.json")]
    policy_cids: PathBuf,
}

impl DeployCommand {
    /// Convert seconds to blocks based on chain's block time
    async fn seconds_to_blocks(rpc_url: &str, seconds: u32) -> Result<u32> {
        let block_time_ms = get_block_time_ms(rpc_url).await?;
        let block_time_seconds = block_time_ms / 1000;

        if block_time_seconds == 0 {
            return Err(eyre::eyre!("Block time is zero, cannot convert seconds to blocks"));
        }

        let blocks = (seconds as u64 / block_time_seconds) as u32;
        tracing::info!(
            "Converting {} seconds to blocks: block_time={}ms ({}s), result={} blocks",
            seconds,
            block_time_ms,
            block_time_seconds,
            blocks
        );

        Ok(blocks)
    }

    /// Get task generator address based on chain ID
    ///
    /// Returns the default task generator address for the given chain:
    /// - Anvil (31337) or empty: 0x70997970C51812dc3A010C7d01b50e0d17dc79C8
    /// - Mainnet (1): 0x70997970C51812dc3A010C7d01b50e0d17dc79C8
    /// - Testnet (others): 0xD45062003a4626a532F30A4596aB253c45AE0647
    async fn get_task_generator_address(rpc_url: &str) -> Result<Address> {
        let chain_id = get_chain_id(rpc_url).await;
        let task_generator_address = match chain_id {
            31337 | 1 => "0x70997970C51812dc3A010C7d01b50e0d17dc79C8",
            _ => "0xD45062003a4626a532F30A4596aB253c45AE0647", // testnet
        };
        task_generator_address
            .parse()
            .map_err(|e| eyre::eyre!("Failed to parse task generator address: {}", e))
    }

    /// Deploy policy data contract
    ///
    /// This function deploys a policy data contract using the Rust chainio functions.
    /// It returns the deployed policy data address.
    /// Note: expire_after_blocks should be in blocks, not seconds.
    async fn deploy_policy_data(
        private_key: &str,
        rpc_url: &str,
        wasm_cid: &str,
        wasm_args_cid: &str,
        policy_data_metadata_cid: &str,
        expire_after_blocks: u32,
    ) -> Result<Address> {
        let controller = PolicyDataController::new(private_key.to_string(), rpc_url.to_string());

        tracing::info!(
            "Deploying policy data:\n  wasmCid: {}\n  wasmArgs: {}\n  metadataCid: {}\n  expireAfter: {} blocks \n",
            wasm_cid,
            wasm_args_cid,
            policy_data_metadata_cid,
            expire_after_blocks
        );

        let (_receipt, policy_data_address) = controller
            .deploy_policy_data(
                wasm_cid.to_string(),
                wasm_args_cid.to_string(),
                expire_after_blocks,
                policy_data_metadata_cid.to_string(),
            )
            .await
            .map_err(|e| eyre::eyre!("Failed to deploy policy data: {}", e))?;

        tracing::info!("Policy data deployed successfully at address: {}", policy_data_address);

        Ok(policy_data_address)
    }

    /// Set attestation info on deployed policy data
    ///
    /// This function sets the attestation info on a deployed policy data contract.
    /// It creates an AttestationInfo struct with:
    /// - attesters: [attester_address, task_generator_address]
    /// - attestationType: ECDSA (0)
    /// - verifier: address(0)
    /// - verificationKey: bytes32(0)
    async fn set_attestation_info(
        private_key: &str,
        rpc_url: &str,
        policy_data_address: Address,
        attester_address: Address,
    ) -> Result<()> {
        let controller = PolicyDataController::new(private_key.to_string(), rpc_url.to_string());

        // Get task generator address based on chain ID
        let task_generator_address = Self::get_task_generator_address(rpc_url).await?;

        // Create attesters array: [attester, task_generator]
        let attesters = vec![attester_address, task_generator_address];

        // Create AttestationInfo struct
        let attestation_info = INewtonPolicyData::AttestationInfo {
            attesters,
            attestationType: 0u8, // ECDSA
            verifier: Address::ZERO,
            verificationKey: FixedBytes::ZERO,
        };

        tracing::info!(
            "Setting attestation info on policy data {} with attesters: {:?}",
            policy_data_address,
            attestation_info.attesters
        );

        controller
            .set_attestation_info(policy_data_address, attestation_info)
            .await
            .map_err(|e| eyre::eyre!("Failed to set attestation info: {}", e))?;

        tracing::info!(
            "Attestation info set successfully on policy data: {}",
            policy_data_address
        );

        Ok(())
    }

    /// Execute the deploy command
    pub async fn execute(self: Box<Self>, _config: NewtonAvsConfig<NewtonCliConfig>) -> eyre::Result<()> {
        // Get values from args or env, with error if still missing
        let private_key = self
            .private_key
            .ok_or_else(|| eyre::eyre!("private_key is required (use --private-key or PRIVATE_KEY env var)"))?;

        let rpc_url = self
            .rpc_url
            .ok_or_else(|| eyre::eyre!("rpc_url is required (use --rpc-url or RPC_URL env var)"))?;

        // Read from policy_cids.json file
        let json_content = std::fs::read_to_string(&self.policy_cids)
            .with_context(|| format!("Failed to read policy_cids.json: {:?}", self.policy_cids))?;
        let json: Value = serde_json::from_str(&json_content)
            .with_context(|| format!("Failed to parse policy_cids.json: {:?}", self.policy_cids))?;

        // Extract values from policy_cids.json
        let wasm_cid = json.get("wasmCid").and_then(|v| v.as_str()).unwrap_or("");
        let wasm_args_cid = json.get("wasmArgs").and_then(|v| v.as_str()).unwrap_or("");
        let policy_data_metadata_cid = json.get("policyDataMetadataCid").and_then(|v| v.as_str()).unwrap_or("");
        let attester = json.get("attester").and_then(|v| v.as_str()).unwrap_or("");
        let attester_address: Address = attester.parse().unwrap_or(Address::ZERO);

        // Convert 300 seconds to blocks (matching the shell script's uint32(300 seconds))
        // The contract expects expire_after in blocks, not seconds
        let expire_after_blocks = Self::seconds_to_blocks(&rpc_url, 300).await?;

        // Deploy policy data
        let policy_data_address = Self::deploy_policy_data(
            &private_key,
            &rpc_url,
            wasm_cid,
            wasm_args_cid,
            policy_data_metadata_cid,
            expire_after_blocks,
        )
        .await?;

        tracing::info!("Policy data deployment completed! Setting attestation info...");
        // Set attestation info on policy data
        Self::set_attestation_info(&private_key, &rpc_url, policy_data_address, attester_address).await?;

        tracing::info!("Successfully deployed policy data and set attestation");
        tracing::info!("Policy Data Address: {}", policy_data_address);

        Ok(())
    }
}

impl SimulateCommand {
    /// Execute the simulate command
    pub async fn execute(self: Box<Self>, config: NewtonAvsConfig<NewtonCliConfig>) -> eyre::Result<()> {
        info!("Executing WASM simulation...");
        let output = utils::execute_wasm(self.wasm_file, self.input_json, config).await?;
        info!("WASM simulation output: {}", output);
        Ok(())
    }
}

impl PolicyDataCommand {
    /// Execute the policy-data command
    pub async fn execute(self: Box<Self>, config: NewtonAvsConfig<NewtonCliConfig>) -> eyre::Result<()> {
        match self.subcommand {
            PolicyDataSubcommand::Deploy(command) => {
                Box::new(command).execute(config).await?;
            }
            PolicyDataSubcommand::Simulate(command) => {
                Box::new(command).execute(config).await?;
            }
        }
        Ok(())
    }
}
