use crate::{
    call_data,
    order_book::{
        CreateOrderParams,
        OrderBookManager,
    },
    parallel_nonce::{
        generate_parallel_session_digest,
        generate_parallel_session_signing_payload,
    },
    trial_trade_account_deploy::{
        CallContractArg,
        CallParams,
        ParallelCallContractArgs,
        ParallelMultiCallContractArgs,
        ParallelRevokeArgs,
        ParallelSessionArgs,
        Payout,
        Secp256k1,
        Session,
        Signature,
        Time,
        TrialLinkAccountArgs as TrialLinkAccountPayload,
        TrialPayoutArgs as TrialPayoutPayload,
        TrialTradeAccountDeploy,
        TrialTradingAccount,
        TrialTradingAccountProxy,
    },
};
use fuels::{
    accounts::signers::private_key::PrivateKeySigner,
    core::traits::Tokenizable,
    crypto::Message,
    prelude::*,
    types::{
        Bytes32,
        Identity,
        U256,
    },
};
use rand::{
    SeedableRng,
    rngs::StdRng,
};
use sha2::{
    Digest,
    Sha256,
};
use std::time::{
    SystemTime,
    UNIX_EPOCH,
};

const TRADE_APPROVAL_DIGEST_TAG: u8 = 0x02;
const PROXY_OWNER_OPERATION_DIGEST_TAG: u8 = 0x03;
const LINK_DIGEST_TAG: u8 = 0x04;

pub struct TrialCallContractParams {
    pub contract_id: ContractId,
    pub function_selector: Vec<u8>,
    pub forward: CallParams,
    pub args: Option<Vec<u8>>,
    pub nonce: U256,
    pub variable_outputs: Option<u16>,
}

#[derive(Debug, Clone)]
pub enum TrialContractCalls {
    Call(TrialCallContractArgs),
    Calls(TrialCallContractsArgs),
}

impl TrialContractCalls {
    pub fn contract_ids(&self) -> Vec<ContractId> {
        match self {
            TrialContractCalls::Call(call) => {
                vec![call.call.call_contract_args.contract_id]
            }
            TrialContractCalls::Calls(calls) => calls
                .calls
                .call_contract_args
                .iter()
                .map(|c| c.contract_id)
                .collect(),
        }
    }
}

#[derive(Debug, Clone)]
pub struct TrialCallContractArgs {
    pub user_signature: Signature,
    pub cosigner_signature: Signature,
    pub call: ParallelCallContractArgs,
    pub variable_outputs: Option<u16>,
    pub contracts: Option<Vec<ContractId>>,
}

#[derive(Debug, Clone)]
pub struct TrialCallContractsArgs {
    pub user_signature: Signature,
    pub cosigner_signature: Signature,
    pub calls: ParallelMultiCallContractArgs,
    pub variable_outputs: Option<u16>,
    pub contracts: Option<Vec<ContractId>>,
}

#[derive(Debug, Clone)]
pub struct TrialCosignerCallRequest {
    pub signature: Signature,
    pub args: ParallelCallContractArgs,
    pub variable_outputs: Option<u16>,
    pub contracts: Option<Vec<ContractId>>,
}

#[derive(Debug, Clone)]
pub struct TrialCosignerCallContractsRequest {
    pub signature: Signature,
    pub args: ParallelMultiCallContractArgs,
    pub variable_outputs: Option<u16>,
    pub contracts: Option<Vec<ContractId>>,
}

#[derive(Debug, Clone)]
pub struct TrialPayoutRequest {
    pub signature: Signature,
    pub args: TrialPayoutPayload,
    pub contracts: Option<Vec<ContractId>>,
}

#[derive(Debug, Clone)]
pub struct TrialSetSessionRequest {
    pub args: ParallelSessionArgs,
    pub contracts: Option<Vec<ContractId>>,
}

#[derive(Debug, Clone)]
pub struct TrialRevokeSessionRequest {
    pub args: ParallelRevokeArgs,
    pub contracts: Option<Vec<ContractId>>,
}

#[derive(Debug, Clone)]
pub struct TrialLinkAccountRequest {
    pub session_signature: Signature,
    pub args: TrialLinkAccountPayload,
    pub contracts: Option<Vec<ContractId>>,
}

impl TrialCallContractParams {
    pub fn new(
        contract_id: ContractId,
        function_selector: Vec<u8>,
        forward: CallParams,
        args: Option<Vec<u8>>,
        nonce: U256,
        variable_outputs: Option<u16>,
    ) -> Self {
        Self {
            contract_id,
            function_selector,
            forward,
            args,
            nonce,
            variable_outputs,
        }
    }

    pub fn call_args(&self) -> CallContractArg {
        CallContractArg {
            contract_id: self.contract_id,
            function_selector: Bytes(self.function_selector.clone()),
            call_params: self.forward.clone(),
            call_data: self.args.clone().map(Bytes),
        }
    }

    pub fn user_message(&self) -> Message {
        generate_parallel_session_signing_payload(self.nonce, self.call_args())
    }

    pub fn trade_approval_message(
        &self,
        chain_id: u64,
        account_contract_id: ContractId,
    ) -> Message {
        generate_trial_trade_approval_signing_payload(
            chain_id,
            account_contract_id,
            self.nonce,
            self.call_args(),
        )
    }
}

fn sha256_digest(bytes: &[u8]) -> [u8; 32] {
    Sha256::digest(bytes).into()
}

fn message_from_digest(digest: [u8; 32]) -> Message {
    Message::from_bytes(digest)
}

pub fn generate_trial_trade_approval_digest<T: Tokenizable>(
    chain_id: u64,
    account_contract_id: ContractId,
    nonce: U256,
    call_contract_arg: T,
) -> [u8; 32] {
    let session_digest = generate_parallel_session_digest(nonce, call_contract_arg);
    generate_trial_trade_approval_digest_from_session_digest(
        chain_id,
        account_contract_id,
        session_digest,
    )
}

pub fn generate_trial_trade_approval_digest_from_session_digest(
    chain_id: u64,
    account_contract_id: ContractId,
    session_digest: [u8; 32],
) -> [u8; 32] {
    let mut bytes = Vec::new();
    bytes.push(TRADE_APPROVAL_DIGEST_TAG);
    bytes.extend_from_slice(&chain_id.to_be_bytes());
    bytes.extend_from_slice(account_contract_id.as_slice());
    bytes.extend_from_slice(&session_digest);
    sha256_digest(&bytes)
}

pub fn generate_trial_trade_approval_signing_payload<T: Tokenizable>(
    chain_id: u64,
    account_contract_id: ContractId,
    nonce: U256,
    call_contract_arg: T,
) -> Message {
    message_from_digest(generate_trial_trade_approval_digest(
        chain_id,
        account_contract_id,
        nonce,
        call_contract_arg,
    ))
}

pub fn generate_trial_proxy_owner_operation_digest<T: Tokenizable>(
    chain_id: u64,
    account_contract_id: ContractId,
    args: T,
) -> [u8; 32] {
    let mut bytes = Vec::new();
    bytes.push(PROXY_OWNER_OPERATION_DIGEST_TAG);
    bytes.extend_from_slice(&chain_id.to_be_bytes());
    bytes.extend_from_slice(account_contract_id.as_slice());
    bytes.extend_from_slice(&call_data!(args));
    sha256_digest(&bytes)
}

pub fn generate_trial_proxy_owner_operation_signing_payload<T: Tokenizable>(
    chain_id: u64,
    account_contract_id: ContractId,
    args: T,
) -> Message {
    message_from_digest(generate_trial_proxy_owner_operation_digest(
        chain_id,
        account_contract_id,
        args,
    ))
}

pub fn generate_trial_link_account_digest<T: Tokenizable>(
    chain_id: u64,
    account_contract_id: ContractId,
    args: T,
) -> [u8; 32] {
    let mut bytes = Vec::new();
    bytes.push(LINK_DIGEST_TAG);
    bytes.extend_from_slice(&chain_id.to_be_bytes());
    bytes.extend_from_slice(account_contract_id.as_slice());
    bytes.extend_from_slice(&call_data!(args));
    sha256_digest(&bytes)
}

pub fn generate_trial_link_account_signing_payload<T: Tokenizable>(
    chain_id: u64,
    account_contract_id: ContractId,
    args: T,
) -> Message {
    message_from_digest(generate_trial_link_account_digest(
        chain_id,
        account_contract_id,
        args,
    ))
}

#[derive(Clone)]
pub struct TrialTradeAccountManager<Wallet: Account + Clone> {
    pub owner: Wallet,
    pub oracle_id: Option<ContractId>,
    pub proxy: TrialTradingAccountProxy<Wallet>,
    pub contract: TrialTradingAccount<Wallet>,
    pub session_signer: Option<PrivateKeySigner>,
    pub cosigner_signer: Option<PrivateKeySigner>,
}

impl TrialTradeAccountManager<Wallet> {
    pub fn contract_id(&self) -> ContractId {
        self.contract.id()
    }

    pub fn identity(&self) -> Identity {
        Identity::ContractId(self.contract.contract_id())
    }

    pub fn session_signer(&self) -> anyhow::Result<PrivateKeySigner> {
        self.session_signer
            .clone()
            .ok_or_else(|| anyhow::anyhow!("Session signer not initialized"))
    }

    pub fn cosigner_signer(&self) -> anyhow::Result<PrivateKeySigner> {
        self.cosigner_signer
            .clone()
            .ok_or_else(|| anyhow::anyhow!("Cosigner signer not initialized"))
    }

    pub async fn new(owner: &Wallet, contract_id: ContractId) -> anyhow::Result<Self> {
        let proxy = TrialTradingAccountProxy::new(contract_id, owner.clone());
        let contract = TrialTradingAccount::new(contract_id, owner.clone());
        Ok(Self {
            owner: owner.clone(),
            oracle_id: None,
            contract,
            proxy,
            session_signer: None,
            cosigner_signer: None,
        })
    }

    pub fn create(
        owner: &Wallet,
        deployment: &TrialTradeAccountDeploy<Wallet>,
    ) -> anyhow::Result<Self> {
        if deployment.proxy_id.is_none() || deployment.proxy.is_none() {
            return Err(anyhow::anyhow!(
                "Trial trade account deploy proxy ID or instance is not set"
            ));
        }
        let contract =
            TrialTradingAccount::new(deployment.proxy_id.unwrap(), owner.clone());
        Ok(Self {
            owner: owner.clone(),
            contract,
            oracle_id: Some(deployment.oracle_id),
            session_signer: None,
            cosigner_signer: None,
            proxy: deployment.proxy.clone().unwrap(),
        })
    }

    pub fn create_with_signers(
        owner: &Wallet,
        deployment: &TrialTradeAccountDeploy<Wallet>,
        contract_ids: &[ContractId],
        cosigner_signer: PrivateKeySigner,
    ) -> anyhow::Result<Self> {
        let mut trial_account = Self::create(owner, deployment)?;
        let mut rng = StdRng::seed_from_u64(2322u64);
        let session_signer = PrivateKeySigner::random(&mut rng);
        let session_address = Identity::Address(session_signer.address());
        trial_account.session_signer = Some(session_signer);
        trial_account.cosigner_signer = Some(cosigner_signer);

        let _session = trial_account.new_session(session_address, contract_ids, None);

        Ok(trial_account)
    }

    pub fn new_session(
        &self,
        session_address: Identity,
        contract_ids: &[ContractId],
        expiry: Option<u64>,
    ) -> Session {
        let default_expiry = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("Time went backwards")
            .as_millis() as u64
            + (30 * 24 * 60 * 60 * 1000);
        Session {
            session_id: session_address,
            expiry: Time {
                unix: expiry.unwrap_or(default_expiry),
            },
            contract_ids: contract_ids.to_vec(),
        }
    }

    async fn chain_id(&self) -> anyhow::Result<u64> {
        Ok(*self
            .owner
            .provider()
            .consensus_parameters()
            .await?
            .chain_id())
    }

    fn tx_contract_ids(
        &self,
        required: impl IntoIterator<Item = ContractId>,
        extra: Option<Vec<ContractId>>,
    ) -> Vec<ContractId> {
        let mut contract_ids = extra.unwrap_or_default();

        for contract_id in required {
            if !contract_ids.contains(&contract_id) {
                contract_ids.push(contract_id);
            }
        }

        if let Some(oracle_id) = self.oracle_id
            && !contract_ids.contains(&oracle_id)
        {
            contract_ids.push(oracle_id);
        }

        contract_ids
    }

    async fn create_proxy_owner_operation_signature<T: Tokenizable>(
        &self,
        args: T,
    ) -> anyhow::Result<Signature> {
        let chain_id = self.chain_id().await?;
        let message = generate_trial_proxy_owner_operation_signing_payload(
            chain_id,
            self.contract_id(),
            args,
        );
        let signature = self.cosigner_signer()?.sign(message).await?;
        Ok(Signature::Secp256k1(Secp256k1 {
            bits: signature.as_slice().try_into()?,
        }))
    }

    pub async fn create_call_args(
        &self,
        params: TrialCallContractParams,
    ) -> anyhow::Result<TrialCallContractArgs> {
        let chain_id = self.chain_id().await?;
        let call_arg = params.call_args();
        let call = ParallelCallContractArgs {
            nonce: params.nonce,
            call_contract_args: call_arg,
        };
        let user_message = generate_parallel_session_signing_payload(
            call.nonce,
            call.call_contract_args.clone(),
        );
        let trade_approval_message = generate_trial_trade_approval_signing_payload(
            chain_id,
            self.contract_id(),
            call.nonce,
            call.call_contract_args.clone(),
        );
        let user_signature = self.session_signer()?.sign(user_message).await?;
        let cosigner_signature =
            self.cosigner_signer()?.sign(trade_approval_message).await?;

        Ok(TrialCallContractArgs {
            user_signature: Signature::Secp256k1(Secp256k1 {
                bits: user_signature.as_slice().try_into()?,
            }),
            cosigner_signature: Signature::Secp256k1(Secp256k1 {
                bits: cosigner_signature.as_slice().try_into()?,
            }),
            variable_outputs: params.variable_outputs,
            call,
            contracts: None,
        })
    }

    pub async fn create_calls_args(
        &self,
        params: Vec<TrialCallContractParams>,
    ) -> anyhow::Result<TrialCallContractsArgs> {
        let call_contract_args = params
            .iter()
            .map(|param| param.call_args())
            .collect::<Vec<_>>();
        let nonce = params
            .first()
            .map(|param| param.nonce)
            .unwrap_or_else(U256::zero);
        let calls = ParallelMultiCallContractArgs {
            nonce,
            call_contract_args,
        };
        let chain_id = self.chain_id().await?;
        let user_message = generate_parallel_session_signing_payload(
            calls.nonce,
            calls.call_contract_args.clone(),
        );
        let trade_approval_message = generate_trial_trade_approval_signing_payload(
            chain_id,
            self.contract_id(),
            calls.nonce,
            calls.call_contract_args.clone(),
        );
        let user_signature = self.session_signer()?.sign(user_message).await?;
        let cosigner_signature =
            self.cosigner_signer()?.sign(trade_approval_message).await?;

        Ok(TrialCallContractsArgs {
            user_signature: Signature::Secp256k1(Secp256k1 {
                bits: user_signature.as_slice().try_into()?,
            }),
            cosigner_signature: Signature::Secp256k1(Secp256k1 {
                bits: cosigner_signature.as_slice().try_into()?,
            }),
            calls,
            variable_outputs: None,
            contracts: None,
        })
    }

    pub fn call_contract(
        &self,
        call_args: &TrialCallContractArgs,
    ) -> CallHandler<Wallet, fuels::programs::calls::ContractCall, ()> {
        let contracts = self.tx_contract_ids(
            [call_args.call.call_contract_args.contract_id],
            call_args.contracts.clone(),
        );
        self.contract
            .methods()
            .call_contract(
                call_args.user_signature.clone(),
                call_args.cosigner_signature.clone(),
                call_args.call.clone(),
            )
            .with_contract_ids(&contracts)
            .with_variable_output_policy(VariableOutputPolicy::Exactly(
                call_args.variable_outputs.unwrap_or_default() as usize,
            ))
    }

    pub fn call_contracts(
        &self,
        call_args: &TrialCallContractsArgs,
    ) -> CallHandler<Wallet, fuels::programs::calls::ContractCall, ()> {
        let contracts = self.tx_contract_ids(
            call_args
                .calls
                .call_contract_args
                .iter()
                .map(|call| call.contract_id),
            call_args.contracts.clone(),
        );
        self.contract
            .methods()
            .call_contracts(
                call_args.user_signature.clone(),
                call_args.cosigner_signature.clone(),
                call_args.calls.clone(),
            )
            .with_contract_ids(&contracts)
            .with_variable_output_policy(VariableOutputPolicy::Exactly(
                call_args.variable_outputs.unwrap_or_default() as usize,
            ))
    }

    pub async fn create_link_account_args(
        &self,
        nonce: U256,
        parent_account: ContractId,
    ) -> anyhow::Result<TrialLinkAccountRequest> {
        let args = TrialLinkAccountPayload {
            nonce,
            parent_account,
        };
        let message = generate_trial_link_account_signing_payload(
            self.chain_id().await?,
            self.contract_id(),
            args.clone(),
        );
        let session_signature = self.session_signer()?.sign(message).await?;
        Ok(TrialLinkAccountRequest {
            session_signature: Signature::Secp256k1(Secp256k1 {
                bits: session_signature.as_slice().try_into()?,
            }),
            args,
            contracts: None,
        })
    }

    pub fn link_account(
        &self,
        link_args: &TrialLinkAccountRequest,
    ) -> CallHandler<Wallet, fuels::programs::calls::ContractCall, ()> {
        let contracts = self.tx_contract_ids([], link_args.contracts.clone());
        self.contract
            .methods()
            .link_account(link_args.session_signature.clone(), link_args.args.clone())
            .with_contract_ids(&contracts)
    }

    pub async fn create_cosigner_call_args(
        &self,
        params: TrialCallContractParams,
    ) -> anyhow::Result<TrialCosignerCallRequest> {
        let call = params.call_args();
        let args = ParallelCallContractArgs {
            nonce: params.nonce,
            call_contract_args: call,
        };
        let signature = self
            .create_proxy_owner_operation_signature(args.clone())
            .await?;
        Ok(TrialCosignerCallRequest {
            signature,
            args,
            variable_outputs: params.variable_outputs,
            contracts: None,
        })
    }

    pub async fn create_cosigner_calls_args(
        &self,
        params: Vec<TrialCallContractParams>,
    ) -> anyhow::Result<TrialCosignerCallContractsRequest> {
        let calls = params
            .iter()
            .map(|param| param.call_args())
            .collect::<Vec<_>>();
        let nonce = params
            .first()
            .map(|param| param.nonce)
            .unwrap_or_else(U256::zero);
        let args = ParallelMultiCallContractArgs {
            nonce,
            call_contract_args: calls,
        };
        let signature = self
            .create_proxy_owner_operation_signature(args.clone())
            .await?;
        Ok(TrialCosignerCallContractsRequest {
            signature,
            args,
            variable_outputs: None,
            contracts: None,
        })
    }

    pub fn cosigner_call(
        &self,
        call_args: &TrialCosignerCallRequest,
    ) -> CallHandler<Wallet, fuels::programs::calls::ContractCall, ()> {
        let contracts = self.tx_contract_ids(
            [call_args.args.call_contract_args.contract_id],
            call_args.contracts.clone(),
        );
        self.contract
            .methods()
            .cosigner_call(Some(call_args.signature.clone()), call_args.args.clone())
            .with_contract_ids(&contracts)
            .with_variable_output_policy(VariableOutputPolicy::Exactly(
                call_args.variable_outputs.unwrap_or_default() as usize,
            ))
    }

    pub fn cosigner_call_contracts(
        &self,
        call_args: &TrialCosignerCallContractsRequest,
    ) -> CallHandler<Wallet, fuels::programs::calls::ContractCall, ()> {
        let contracts = self.tx_contract_ids(
            call_args
                .args
                .call_contract_args
                .iter()
                .map(|call| call.contract_id),
            call_args.contracts.clone(),
        );
        self.contract
            .methods()
            .cosigner_call_contracts(
                Some(call_args.signature.clone()),
                call_args.args.clone(),
            )
            .with_contract_ids(&contracts)
            .with_variable_output_policy(VariableOutputPolicy::Exactly(
                call_args.variable_outputs.unwrap_or_default() as usize,
            ))
    }

    pub async fn create_payout_args(
        &self,
        nonce: U256,
        payouts: Vec<Payout>,
    ) -> anyhow::Result<TrialPayoutRequest> {
        let args = TrialPayoutPayload { nonce, payouts };
        let signature = self
            .create_proxy_owner_operation_signature(args.clone())
            .await?;
        Ok(TrialPayoutRequest {
            signature,
            args,
            contracts: None,
        })
    }

    pub fn payout(
        &self,
        payout_args: &TrialPayoutRequest,
    ) -> CallHandler<Wallet, fuels::programs::calls::ContractCall, ()> {
        let contracts = self.tx_contract_ids([], payout_args.contracts.clone());
        self.contract
            .methods()
            .payout(
                Some(payout_args.signature.clone()),
                payout_args.args.clone(),
            )
            .with_contract_ids(&contracts)
    }

    pub fn create_set_session_args(
        &self,
        nonce: U256,
        session: Session,
    ) -> TrialSetSessionRequest {
        let args = ParallelSessionArgs {
            nonce,
            session_id: session.session_id,
            expiry: session.expiry,
            contract_ids: session.contract_ids,
        };
        TrialSetSessionRequest {
            args,
            contracts: None,
        }
    }

    pub fn set_session(
        &self,
        session_args: &TrialSetSessionRequest,
    ) -> CallHandler<Wallet, fuels::programs::calls::ContractCall, ()> {
        let contracts = self.tx_contract_ids([], session_args.contracts.clone());
        self.contract
            .methods()
            .set_session(session_args.args.clone())
            .with_contract_ids(&contracts)
    }

    pub fn create_revoke_session_args(&self, nonce: U256) -> TrialRevokeSessionRequest {
        let args = ParallelRevokeArgs { nonce };
        TrialRevokeSessionRequest {
            args,
            contracts: None,
        }
    }

    pub fn revoke_session(
        &self,
        revoke_args: &TrialRevokeSessionRequest,
    ) -> CallHandler<Wallet, fuels::programs::calls::ContractCall, ()> {
        let contracts = self.tx_contract_ids([], revoke_args.contracts.clone());
        self.contract
            .methods()
            .revoke_session(revoke_args.args.clone())
            .with_contract_ids(&contracts)
    }

    pub fn cancel_orders_args(
        &self,
        order_book: &OrderBookManager<Wallet>,
        order_id: Bytes32,
        gas: Option<u64>,
        nonce: U256,
    ) -> TrialCallContractParams {
        TrialCallContractParams::new(
            order_book.contract.contract_id(),
            crate::fn_selector!(cancel_order(OrderId)),
            CallParams {
                coins: 0,
                asset_id: AssetId::default(),
                gas: gas.unwrap_or(u64::MAX),
            },
            Some(call_data!(*order_id)),
            nonce,
            None,
        )
    }

    pub async fn cancel_order(
        &mut self,
        order_book: &OrderBookManager<Wallet>,
        order_id: Bytes32,
        gas: Option<u64>,
        nonce: U256,
    ) -> anyhow::Result<TrialCallContractArgs> {
        self.create_call_args(self.cancel_orders_args(order_book, order_id, gas, nonce))
            .await
    }

    pub fn create_order_args(
        &self,
        order_book: &OrderBookManager<Wallet>,
        params: &CreateOrderParams,
        gas: Option<u64>,
        nonce: U256,
    ) -> TrialCallContractParams {
        let forward = order_book.create_call_params(params, gas);
        TrialCallContractParams::new(
            order_book.contract.contract_id(),
            crate::fn_selector!(create_order(OrderArgs)),
            CallParams {
                coins: forward.coins,
                asset_id: forward.asset_id,
                gas: forward.gas,
            },
            Some(call_data!(params.to_order_args())),
            nonce,
            None,
        )
    }

    pub async fn create_order(
        &mut self,
        order_book: &OrderBookManager<Wallet>,
        params: &CreateOrderParams,
        gas: Option<u64>,
        nonce: U256,
    ) -> anyhow::Result<TrialCallContractArgs> {
        self.create_call_args(self.create_order_args(order_book, params, gas, nonce))
            .await
    }
}
