mod evm;
mod near;
mod svm;

use alloy_consensus::{TxEip1559, TxEip2930, TxEip4844, TxEip7702};
use alloy_dyn_abi::JsonAbiExt;
use alloy_json_abi::JsonAbi;
use alloy_rlp::Decodable;
use anchor_idl::Idl;
use base64::engine::general_purpose;
use base64::Engine;
use borsh::BorshDeserialize;
use dew_schema_language::engine::{DewSchemaLanguageEngine, DewSchemaLanguageResult};
use near_sdk::{log, near, AccountId};
use serde_json::{json, Value};
use solana_message::{compiled_instruction::CompiledInstruction, VersionedMessage};
use std::collections::{HashMap, HashSet};

pub use near::Action;

use svm::TokenInstruction;

use crate::{
    chainsig::get_network_by_account_id,
    helpers::{
        chainsig::{derived_public_key_with_address, ChainEnvironment},
        decoder::{evm::decode_dynsol, near::Transaction, svm::decode_instruction},
    },
    permission::types::{NearNativeTransactionConfig, Restriction},
};

#[derive(Debug)]
#[near(serializers = [json])]
pub struct TransactionInput {
    contract_id: String,
    function_name: String,
    args: Value,
}

pub enum UnsignedTx {
    Eip2930(TxEip2930),
    Eip1559(TxEip1559),
    Eip4844(TxEip4844),
    Eip7702(TxEip7702),
}

pub fn decode_unsigned(raw: &[u8]) -> UnsignedTx {
    match raw[0] {
        0x01 => UnsignedTx::Eip2930(TxEip2930::decode(&mut &raw[1..]).unwrap()),
        0x02 => UnsignedTx::Eip1559(TxEip1559::decode(&mut &raw[1..]).unwrap()),
        0x03 => UnsignedTx::Eip4844(TxEip4844::decode(&mut &raw[1..]).unwrap()),
        0x04 => UnsignedTx::Eip7702(TxEip7702::decode(&mut &raw[1..]).unwrap()),
        _ => panic!("Legacy transactions are not supported"),
    }
}

pub fn parse_transaction(
    encoded_transaction: &String,
    strategy_config: &NearNativeTransactionConfig,
) -> Vec<TransactionInput> {
    let chain = &strategy_config.chain_environment;

    let transaction_inputs: Option<Vec<TransactionInput>> = if let ChainEnvironment::EVM = chain {
        let first_schema = strategy_config
            .restrictions
            .first()
            .expect("Missing schema");
        let interface_base64 = &first_schema.interface;

        let decoded_bytes = general_purpose::STANDARD
            .decode(interface_base64.as_bytes())
            .expect("Failed to decode interface");

        let interface =
            String::from_utf8(decoded_bytes).expect("Failed to convert bytes to string");

        let raw_bytes = hex::decode(encoded_transaction.trim_start_matches("0x")).unwrap();

        let parsed_transaction = decode_unsigned(&raw_bytes);

        let (selector, args_data, receiver_id) = match parsed_transaction {
            UnsignedTx::Eip2930(tx) => (
                &tx.input.clone()[0..4],
                &tx.input.clone()[4..],
                tx.to.to().unwrap().to_string(),
            ),
            UnsignedTx::Eip1559(tx) => (
                &tx.input.clone()[0..4],
                &tx.input.clone()[4..],
                tx.to.to().unwrap().to_string(),
            ),
            UnsignedTx::Eip4844(tx) => (
                &tx.input.clone()[0..4],
                &tx.input.clone()[4..],
                tx.to.to_string(),
            ),
            UnsignedTx::Eip7702(tx) => (
                &tx.input.clone()[0..4],
                &tx.input.clone()[4..],
                tx.to.to_string(),
            ),
        };

        let abi: JsonAbi = serde_json::from_str(&interface).unwrap();

        let func = abi
            .functions()
            .find(|f| f.selector() == selector)
            .expect("Function not found in ABI");

        let decoded_input = func.abi_decode_input(args_data).unwrap();
        let mut args_json = serde_json::Map::new();
        for (param, value) in func.inputs.iter().zip(decoded_input.iter()) {
            let arg = decode_dynsol(value);
            args_json.insert(param.name.clone(), arg);
        }
        let input = TransactionInput {
            contract_id: receiver_id,
            function_name: func.name.to_owned(),
            args: Value::Object(args_json),
        };

        vec![input].into()
    } else if let ChainEnvironment::SVM = chain {
        let mut all_interfaces: Vec<Idl> = Vec::new();

        for pattern in &strategy_config.restrictions {
            let interface_base64 = &pattern.interface;

            if interface_base64 != "" {
                let decoded_bytes = general_purpose::STANDARD
                    .decode(interface_base64.as_bytes())
                    .expect("Failed to decode interface");

                let interface =
                    String::from_utf8(decoded_bytes).expect("Failed to convert bytes to string");

                let idl: Idl = serde_json::from_str(&interface).unwrap();
                all_interfaces.push(idl);
            }
        }

        let raw_bytes = general_purpose::STANDARD
            .decode(encoded_transaction.as_bytes())
            .expect("Failed to decode message");

        let versioned_msg: VersionedMessage =
            bincode::deserialize(&raw_bytes).expect("failed to decode");

        let mut relevant_instructions: Vec<CompiledInstruction> = Vec::new();
        let instructions = versioned_msg.instructions();
        let irrelevant_program_ids = HashSet::from([
            // Memo program
            "MemoSq4gqABAXKb96qnH8TysNcWxMyWCqXgDLGmfcHr".to_owned(),
            // Compute budget program
            "ComputeBudget111111111111111111111111111111".to_owned(),
        ]);

        let mut transactions_input: Vec<TransactionInput> = vec![];

        instructions.iter().for_each(|instruction| {
            let program_id =
                versioned_msg.static_account_keys()[instruction.program_id_index as usize];

            if irrelevant_program_ids.contains(&program_id.to_string()) == false {
                relevant_instructions.push(instruction.clone());
            }
        });

        let account_keys = versioned_msg.static_account_keys();

        relevant_instructions
            .iter()
            .for_each(|compiled_instruction| {
                let program_id = account_keys[compiled_instruction.program_id_index as usize];

                let possible_idl = all_interfaces
                    .iter()
                    .filter(|idl| idl.address == program_id.to_string());
                let mut is_instruction_parsed = false;
                if program_id.to_string() == "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA" {
                    let instruction = TokenInstruction::unpack(&compiled_instruction.data);

                    match instruction {
                        TokenInstruction::Transfer { amount } => {
                            let source =
                                account_keys[compiled_instruction.accounts[0] as usize].to_string();
                            let destination =
                                account_keys[compiled_instruction.accounts[1] as usize].to_string();
                            let authority =
                                account_keys[compiled_instruction.accounts[2] as usize].to_string();

                            transactions_input.push(TransactionInput {
                                args: json!({
                                    "args": {
                                        "amount": amount.to_string()
                                    },
                                    "accounts": {
                                        "source": source,
                                        "destination": destination,
                                        "authority": authority,
                                    }
                                }),
                                contract_id: program_id.to_string(),
                                function_name: "transfer".to_owned(),
                            })
                        }
                        TokenInstruction::Approve { amount } => {
                            let source =
                                account_keys[compiled_instruction.accounts[0] as usize].to_string();
                            let delegate =
                                account_keys[compiled_instruction.accounts[1] as usize].to_string();
                            let authority =
                                account_keys[compiled_instruction.accounts[2] as usize].to_string();

                            transactions_input.push(TransactionInput {
                                args: json!({
                                    "args": {
                                        "amount": amount.to_string()
                                    },
                                    "accounts": {
                                        "source": source,
                                        "delegate": delegate,
                                        "authority": authority,
                                    }
                                }),
                                contract_id: program_id.to_string(),
                                function_name: "approve".to_owned(),
                            })
                        }
                        TokenInstruction::Revoke => {
                            let source =
                                account_keys[compiled_instruction.accounts[0] as usize].to_string();
                            let authority =
                                account_keys[compiled_instruction.accounts[1] as usize].to_string();

                            transactions_input.push(TransactionInput {
                                args: json!({
                                    "accounts": {
                                        "source": source,
                                        "authority": authority,
                                    }
                                }),
                                contract_id: program_id.to_string(),
                                function_name: "revoke".to_owned(),
                            })
                        }
                        TokenInstruction::TransferChecked { amount, decimals } => {
                            let source =
                                account_keys[compiled_instruction.accounts[0] as usize].to_string();
                            let mint =
                                account_keys[compiled_instruction.accounts[1] as usize].to_string();
                            let destination =
                                account_keys[compiled_instruction.accounts[2] as usize].to_string();
                            let authority =
                                account_keys[compiled_instruction.accounts[3] as usize].to_string();
                            transactions_input.push(TransactionInput {
                                args: json!({
                                    "args": {
                                        "amount": amount.to_string(),
                                        "decimals": decimals.to_string()
                                    },
                                    "accounts": {
                                        "source": source,
                                        "destination": destination,
                                        "authority": authority,
                                        "mint": mint
                                    }
                                }),
                                contract_id: program_id.to_string(),
                                function_name: "transferChecked".to_owned(),
                            })
                        }
                        TokenInstruction::ApproveChecked { amount, decimals } => {
                            let source =
                                account_keys[compiled_instruction.accounts[0] as usize].to_string();
                            let mint =
                                account_keys[compiled_instruction.accounts[1] as usize].to_string();
                            let delegate =
                                account_keys[compiled_instruction.accounts[2] as usize].to_string();
                            let authority =
                                account_keys[compiled_instruction.accounts[3] as usize].to_string();
                            transactions_input.push(TransactionInput {
                                args: json!({
                                    "args": {
                                        "amount": amount.to_string(),
                                        "decimals": decimals.to_string()
                                    },
                                    "accounts": {
                                        "source": source,
                                        "delegate": delegate,
                                        "authority": authority,
                                        "mint": mint
                                    }
                                }),
                                contract_id: program_id.to_string(),
                                function_name: "approveChecked".to_owned(),
                            })
                        }
                    }

                    is_instruction_parsed = true;
                } else {
                    for idl in possible_idl {
                        let instruction = decode_instruction(&idl, &compiled_instruction.data);
                        match instruction {
                            Some(instruction) => {
                                let mut accounts_map = serde_json::Map::new();

                                for (i, acc_index) in
                                    compiled_instruction.accounts.iter().enumerate()
                                {
                                    let key = account_keys[*acc_index as usize];
                                    // we assume it always have 1 instruction only
                                    let label = idl.instructions[0].accounts.get(i).unwrap();
                                    match label {
                                        anchor_idl::IdlInstructionAccountItem::Composite(
                                            _idl_instruction_accounts,
                                        ) => {
                                            panic!("Composite account not supported");
                                        }
                                        anchor_idl::IdlInstructionAccountItem::Single(
                                            idl_instruction_account,
                                        ) => {
                                            accounts_map.insert(
                                                idl_instruction_account.name.clone(),
                                                serde_json::Value::String(key.to_string()),
                                            );
                                        }
                                    }
                                }

                                transactions_input.push(TransactionInput {
                                    args: json!({
                                        "args": instruction.arguments,
                                        "accounts": Value::Object(accounts_map)
                                    }),
                                    contract_id: program_id.to_string(),
                                    function_name: instruction.instruction_name,
                                });
                                is_instruction_parsed = true;
                                break;
                            }
                            None => {}
                        }
                    }
                }

                if is_instruction_parsed == false {
                    panic!("One of the instruction is not parsable, please double check the idl");
                }
            });

        transactions_input.into()
    } else if let ChainEnvironment::NearWasm = chain {
        let tx = parse_near_transaction(encoded_transaction);

        let mut transaction_inputs: Vec<TransactionInput> = Vec::new();

        tx.actions().iter().for_each(|action| {
            // we only parse FunctionCall
            match action {
                Action::FunctionCall(fc) => {
                    let args_string = String::from_utf8(fc.args.clone()).unwrap();
                    let args = serde_json::from_str::<Value>(&args_string).unwrap();

                    transaction_inputs.push(TransactionInput {
                        args,
                        contract_id: tx.receiver_id().to_string(),
                        function_name: fc.method_name.clone(),
                    });
                }
                _ => {}
            }
        });

        transaction_inputs.into()
    } else {
        None
    };

    let input = transaction_inputs.expect("Failed to parse transaction");

    input
}

pub fn parse_near_transaction(encoded_tx: &String) -> Transaction {
    let raw_bytes = general_purpose::STANDARD
        .decode(encoded_tx.as_bytes())
        .expect("Failed to decode message");

    Transaction::try_from_slice(&raw_bytes).expect("Failed to decode NEAR transaction")
}

trait DslResultExt {
    fn as_string(&self) -> &String;
}

impl DslResultExt for DewSchemaLanguageResult {
    fn as_string(&self) -> &String {
        if let DewSchemaLanguageResult::String(v) = self {
            v
        } else {
            panic!("Expected VariantA, got something else");
        }
    }
}

pub fn validate_input(
    transaction_inputs: &Vec<TransactionInput>,
    restrictions: &Vec<Restriction>,
    current_account_id: AccountId,
) -> Result<bool, String> {
    let mut schema_index = 0;
    let mut input_index = 0;

    while schema_index < restrictions.len() {
        let pattern = &restrictions[schema_index];

        let input = transaction_inputs
            .get(input_index)
            .expect("Input is less than schema");

        let mut host_functions: HashMap<
            String,
            Box<
                dyn Fn(
                    Vec<DewSchemaLanguageResult>,
                    Option<&DewSchemaLanguageResult>,
                ) -> Result<DewSchemaLanguageResult, String>,
            >,
        > = HashMap::new();

        // chain_sig_address(path, "SVM" | "EVM" | "NearWasm")
        host_functions.insert(
            "chain_sig_address".into(),
            Box::new({
                let current_account_id = current_account_id.clone();
                move |args, callee| {
                    if callee.is_some() {
                        return Err("chain_sig_address should not be called with a callee".into());
                    }

                    if args.len() != 2 {
                        return Err(
                        "chain_sig_address must have two arguments. chain_sig_address(path, env)"
                            .into(),
                        );
                    }

                    let path = args.get(0).unwrap().as_string();
                    let env: &String = args.get(1).unwrap().as_string();

                    let public_key_address = derived_public_key_with_address(
                        path,
                        &ChainEnvironment::from_string(env.to_string()),
                        &get_network_by_account_id(&current_account_id),
                        &current_account_id,
                    );

                    Ok(DewSchemaLanguageResult::String(public_key_address.address))
                }
            }),
        );

        let root_obj = json!({
            "function_name": input.function_name,
            "contract_id": input.contract_id,
            "args": input.args,
        });

        let engine = DewSchemaLanguageEngine::new(root_obj.to_string(), host_functions);

        let evaluate_result_wrapped = engine.evaluate(pattern.schema.clone()).unwrap();

        if evaluate_result_wrapped == DewSchemaLanguageResult::Boolean(true) {
            schema_index += 1;
            input_index += 1;
        } else {
            match pattern.go_to_index_if_not_found {
                Some(next_schema_index) => {
                    log!(
                        "Schema [{}] failed, skipping to Schema [{}]",
                        schema_index,
                        next_schema_index
                    );
                    schema_index = next_schema_index as usize;
                }
                None => {
                    // cant match and this is a required action
                    // we fail this
                    return Err(format!(
                        "Fail to validate action at schema[{}]",
                        schema_index
                    ));
                }
            }
        }
    }

    if input_index != transaction_inputs.len() {
        return Err("Input is more than schema".to_string());
    }

    Ok(true)
}

pub fn parse_and_validate_transaction(
    encoded_transaction: &String,
    strategy_config: &NearNativeTransactionConfig,
    current_account_id: AccountId,
) -> bool {
    let transaction_inputs = parse_transaction(encoded_transaction, strategy_config);

    let transaction_runnable_result = validate_input(
        &transaction_inputs,
        &strategy_config.restrictions,
        current_account_id,
    );

    if transaction_runnable_result.as_ref().is_err() == true {
        panic!("{}", transaction_runnable_result.as_ref().unwrap_err());
    }

    transaction_runnable_result.unwrap()
}

#[cfg(all(test, not(target_arch = "wasm32")))]
mod test_decoder {
    use near_sdk::log;
    use serde_json::json;

    use crate::{
        helpers::decoder::{validate_input, TransactionInput},
        permission::Restriction,
    };

    #[test]
    fn test_basic_input_with_dsl_valid() {
        let restrictions = vec![Restriction {
            interface: "".to_owned(),
            schema: r##"$.contract_id.equal("rek.near")"##.to_owned(),
            go_to_index_if_not_found: None,
        }];

        let transaction_inputs = vec![TransactionInput {
            contract_id: "rek.near".to_owned(),
            function_name: "redeem".to_owned(),
            args: json!({}),
        }];

        let transaction_runnable_result = validate_input(
            &transaction_inputs,
            &restrictions,
            "alice.near".parse().unwrap(),
        );

        assert_eq!(transaction_runnable_result.as_ref().is_err(), false);
    }

    #[test]
    fn test_basic_input_with_dsl_invalid() {
        let restrictions = vec![Restriction {
            interface: "".to_owned(),
            schema: r##"$.contract_id.equal("rek.near")"##.to_owned(),
            go_to_index_if_not_found: None,
        }];

        let transaction_inputs = vec![TransactionInput {
            contract_id: "alice.near".to_owned(),
            function_name: "redeem".to_owned(),
            args: json!({}),
        }];

        let transaction_runnable_result = validate_input(
            &transaction_inputs,
            &restrictions,
            "alice.near".parse().unwrap(),
        );
        log!("{:?}", transaction_runnable_result);
        log!("{:?}", transaction_runnable_result.as_ref().is_err());
        assert_eq!(transaction_runnable_result.as_ref().is_err(), true);
    }

    #[test]
    fn test_full_input_with_dsl_valid() {
        let restrictions = vec![
            Restriction {
                interface: "".to_owned(),
                // schema: r##"$.args.amount.equal("1000")"##.to_owned(),
                // schema: r##"$.contract_id.equal("alice.near")"##.to_owned(),
                // schema: r##"$.function_name.equal("redeem")"##.to_owned(),
                schema: r##"and($.contract_id.equal("alice.near"),$.function_name.equal("redeem"),$.args.amount.equal("1000"))"##.to_owned(),
                go_to_index_if_not_found: None,
            }
        ];

        let transaction_inputs = vec![TransactionInput {
            contract_id: "alice.near".to_owned(),
            function_name: "redeem".to_owned(),
            args: json!({
                "amount": "1000"
            }),
        }];

        let transaction_runnable_result = validate_input(
            &transaction_inputs,
            &restrictions,
            "alice.near".parse().unwrap(),
        );
        log!("{:?}", transaction_runnable_result);
        log!("{:?}", transaction_runnable_result.as_ref().is_err());
        assert_eq!(transaction_runnable_result.as_ref().is_err(), false);
    }

    #[test]
    fn test_full_input_with_dsl_invalid() {
        let restrictions = vec![
            Restriction {
                interface: "".to_owned(),
                // schema: r##"$.args.amount.equal("1000")"##.to_owned(),
                // schema: r##"$.contract_id.equal("alice.near")"##.to_owned(),
                // schema: r##"$.function_name.equal("redeem")"##.to_owned(),
                schema: r##"and($.contract_id.equal("alicee.near"),$.function_name.equal("redeem"),$.args.amount.equal("1000"))"##.to_owned(),
                go_to_index_if_not_found: None,
            }
        ];

        let transaction_inputs = vec![TransactionInput {
            contract_id: "alice.near".to_owned(),
            function_name: "redeem".to_owned(),
            args: json!({
                "amount": "1000"
            }),
        }];

        let transaction_runnable_result = validate_input(
            &transaction_inputs,
            &restrictions,
            "alice.near".parse().unwrap(),
        );
        log!("{:?}", transaction_runnable_result);
        log!("{:?}", transaction_runnable_result.as_ref().is_err());
        assert_eq!(transaction_runnable_result.as_ref().is_err(), true);
    }

    #[test]
    fn test_chain_sig_address_host_function() {
        let restrictions = vec![Restriction {
            interface: "".to_owned(),
            schema: r##"$.args.destination.case_insensitive_equal(chain_sig_address("2", "EVM"))"##
                .to_owned(),
            go_to_index_if_not_found: None,
        }];

        let transaction_inputs = vec![TransactionInput {
            contract_id: "alice.near".to_owned(),
            function_name: "redeem".to_owned(),
            args: json!({
                "destination": "0x3293a35b76aa87DDaaca85224a32E48Cd091D91a"
            }),
        }];

        let transaction_runnable_result = validate_input(
            &transaction_inputs,
            &restrictions,
            "alice.near".parse().unwrap(),
        );
        log!("{:?}", transaction_runnable_result);
        log!("{:?}", transaction_runnable_result.as_ref().is_err());
        assert_eq!(transaction_runnable_result.as_ref().is_err(), false);
    }

    #[test]
    fn test_json_schema() {
        let restrictions = vec![Restriction {
            interface: "".to_owned(),
            schema: r##"and(
                $.args.msg.json().Execute.action.IncreaseCollateral.token_id.equal("lst.rhealab.near"),
                $.args.msg.json().Execute.action.IncreaseCollateral.max_amount.equal("123")
            )"##
                .to_owned(),
            go_to_index_if_not_found: None,
        }];

        let transaction_inputs = vec![TransactionInput {
            contract_id: "alice.near".to_owned(),
            function_name: "redeem".to_owned(),
            args: json!({
                "msg": r##"{"Execute":{"action":{"IncreaseCollateral":{"token_id":"lst.rhealab.near","max_amount":"123"}}}}"##
            }),
        }];

        let transaction_runnable_result = validate_input(
            &transaction_inputs,
            &restrictions,
            "alice.near".parse().unwrap(),
        );
        log!("{:?}", transaction_runnable_result);
        log!("{:?}", transaction_runnable_result.as_ref().is_err());
        assert_eq!(transaction_runnable_result.as_ref().is_err(), false);
    }
}