mod agent;
mod chainsig;
mod helpers;
mod internal;
mod permission;

use crate::agent::Agent;
use crate::permission::{EmergencyConfig, Policy, PolicyType, Proposal, Role, RoleTarget, Vote};
pub use common_data::asset::Asset;
use common_data::constants::{
    DEFAULT_PROPOSAL_EXPIRY_TIME_NS, NANOSECONDS_PER_DAY, NANOSECONDS_PER_WEEK,
};
use near_sdk::env::block_timestamp;
use near_sdk::json_types::U128;
use near_sdk::store::{IterableMap, IterableSet, LazyOption, LookupMap};
use near_sdk::{assert_one_yocto, env, AccountId};
use near_sdk::{near, require};

pub type ChangeControl = common_data::change_control::ChangeControl;

/// Storage keys for collections to prevent storage collisions
#[derive(near_sdk::BorshStorageKey, borsh::BorshSerialize)]
pub enum StorageKey {
    RoleAssignments,
    Policies,
    Proposals,
    ProposalVotes,
    PendingPolicies,
    PendingActions,
    RegisteredAgents,
    DataStorage,
    EmergencyAccounts,
    EmergencyVotes,
    RoleSet { target_hash: Vec<u8> },
    StagedWasm,
}

#[near(contract_state)]
pub struct Contract {
    // Nano timestamp of when vault is created
    pub created_at: U128,
    // Policy change control configuration
    pub policy_change_control: ChangeControl,
    // Emergency config, immutable after initialization
    pub emergency_config: EmergencyConfig,
    // State of the vault
    kernel_state: KernelState,
    // Staged contract code for upgrades
    pub staged_wasm: LazyOption<Vec<u8>>,
    // Base58 hash of staged contract code
    pub staged_code_hash: Option<String>,
}

/// Kernel state structure containing all contract data
#[near(serializers=[borsh])]
pub struct KernelState {
    // -------------------- Permission System --------------------
    /// RoleTarget -> Roles (supports both AccountId and Codehash targets)
    role_assignments: IterableMap<RoleTarget, IterableSet<Role>>,
    /// PolicyId -> Policy
    policies: IterableMap<String, Policy>,

    // -------------------- Proposal System --------------------
    /// Proposal ID -> Proposal
    proposals: IterableMap<u64, Proposal>,
    /// Proposal ID -> Vec<Vote>
    proposal_votes: LookupMap<u64, Vec<Vote>>,
    /// Counter for generating unique proposal IDs
    proposal_counter: u64,
    /// Last executed proposal timestamp nano
    last_executed_proposal_timestamp: U128,

    // -------------------- Pending Changes --------------------
    /// Pending policy changes by policy_id
    pending_policies: IterableMap<String, Policy>,

    // -------------------- Lock State --------------------
    /// Global lock holder - tracks who currently holds the contract lock
    global_lock_holder: Option<AccountId>,
    /// Pending actions that must be completed before lock release (stores policy IDs)
    pending_actions: IterableSet<String>,

    // -------------------- Agent Registry --------------------
    /// AccountId -> Agent mapping for registered agents
    registered_agent_by_account_id: IterableMap<AccountId, Agent>,

    // -------------------- Data Storage --------------------
    /// String key -> String value data storage
    data_storage: LookupMap<String, String>,

    // -------------------- Emergency Voting System --------------------
    /// Emergency vote ID -> Vec<Vote>
    emergency_votes: IterableMap<AccountId, Vote>,
    /// Indicates that the contract has emergency mode active
    emergency_mode_active: bool,
}

// Implement the contract structure
#[near]
impl Contract {
    /// Check if a pending policy is ready to activate
    pub fn is_policy_ready_to_activate(&self, policy_id: String) -> bool {
        if let Some(pending) = self.kernel_state.pending_policies.get(&policy_id) {
            (near_sdk::env::block_timestamp() as u128) >= pending.activation_time.0
        } else {
            false
        }
    }
}

impl Default for KernelState {
    fn default() -> Self {
        Self {
            role_assignments: IterableMap::new(StorageKey::RoleAssignments),
            policies: IterableMap::new(StorageKey::Policies),
            proposals: IterableMap::new(StorageKey::Proposals),
            proposal_votes: LookupMap::new(StorageKey::ProposalVotes),
            proposal_counter: 0,
            last_executed_proposal_timestamp: U128(block_timestamp() as u128),
            pending_policies: IterableMap::new(StorageKey::PendingPolicies),
            global_lock_holder: None,
            pending_actions: IterableSet::new(StorageKey::PendingActions),
            registered_agent_by_account_id: IterableMap::new(StorageKey::RegisteredAgents),
            data_storage: LookupMap::new(StorageKey::DataStorage),
            emergency_votes: IterableMap::new(StorageKey::EmergencyVotes),
            emergency_mode_active: false,
        }
    }
}

impl Default for Contract {
    fn default() -> Self {
        // Use minimum valid emergency inactive duration (1 week + 1 nanosecond)
        Self::new(
            ChangeControl::default(),
            None,
            U128(NANOSECONDS_PER_WEEK * 2),
            vec![],
            0,
            U128(NANOSECONDS_PER_WEEK),
        )
    }
}

#[near]
impl Contract {
    /// Initialize anew contract instance
    #[init]
    pub fn new(
        policy_change_control: ChangeControl,
        owner: Option<AccountId>,
        emergency_inactive_duration_ns: U128,
        emergency_accounts: Vec<AccountId>,
        emergency_vote_threshold: u32,
        emergency_vote_duration_ns: U128,
    ) -> Self {
        require!(
            emergency_inactive_duration_ns.0 > NANOSECONDS_PER_WEEK,
            "Emergency inactive duration must be greater than one week"
        );
        require!(
            emergency_vote_duration_ns.0 > NANOSECONDS_PER_DAY,
            "Emergency vote duration must be greater than one day"
        );
        require!(
            emergency_vote_duration_ns.0 < emergency_inactive_duration_ns.0,
            "Emergency vote duration must be less than emergency inactive duration"
        );
        require!(
            emergency_vote_threshold <= emergency_accounts.len() as u32,
            "Emergency vote threshold cannot exceed number of emergency accounts"
        );

        let current_time = near_sdk::env::block_timestamp() as u128;

        let mut contract = Self {
            created_at: U128(current_time),
            policy_change_control,
            emergency_config: EmergencyConfig {
                emergency_accounts: IterableSet::new(StorageKey::EmergencyAccounts),
                emergency_inactive_duration_ns,
                emergency_vote_threshold,
                emergency_vote_duration_ns,
            },
            kernel_state: KernelState::default(),
            staged_wasm: LazyOption::new(StorageKey::StagedWasm, None),
            staged_code_hash: None,
        };

        for account in &emergency_accounts {
            contract
                .emergency_config
                .emergency_accounts
                .insert(account.clone());
        }

        // Grant owner role to specified owner or deployer
        let owner_account = owner.unwrap_or_else(|| near_sdk::env::predecessor_account_id());
        let target = RoleTarget::AccountId(owner_account.clone());
        contract.internal_grant_role("owner".to_string(), target);

        contract.internal_initialize_vault();

        contract
    }

    #[init(ignore_state)]
    #[payable]
    pub fn reset_with_config(
        policy_change_control: ChangeControl,
        owner: Option<AccountId>,
    ) -> Self {
        let mut old_contract_mut: Contract = env::state_read().unwrap();

        if env::predecessor_account_id() == env::current_account_id() {
            // This is not typo
            // signer_account_id is used to ensure that only the key holder can reset the state
            // we also need to ensure that caller pay at least 1 yoctoNEAR to ensure that they key is actually a full access key
            assert!(
                env::signer_account_id() == env::current_account_id(),
                "This method must be called directly by the contract account, delegation is not allowed"
            );
            assert_one_yocto();
        } else {
            panic!("Only owner or contract account can reset the state");
        }

        let current_time = near_sdk::env::block_timestamp() as u128;

        old_contract_mut.created_at = U128(current_time);

        old_contract_mut.policy_change_control = policy_change_control;

        if let Some(owner_account) = owner {
            let target = RoleTarget::AccountId(owner_account.clone());
            old_contract_mut.internal_grant_role("owner".to_string(), target);
        }

        // Kernel state should not be reset
        // Contract should not be re-initialized
        old_contract_mut.staged_wasm.set(None);
        old_contract_mut.staged_code_hash = None;

        old_contract_mut
    }

    pub fn upload_contract(&mut self) -> String {
        require!(
            self.internal_has_role(&env::predecessor_account_id(), "owner"),
            "Only owner can upload contract"
        );
        require!(
            self.staged_code_hash.is_none(),
            "Previous contract code hash is not deleted yet"
        );

        let code = env::input().expect("Error: No input");
        require!(!code.is_empty(), "Error: Empty contract code");

        self.staged_wasm.set(Some(code.clone()));

        let hash = bs58::encode(env::sha256(&code)).into_string();
        self.staged_code_hash = Some(hash.clone());

        hash
    }
}

#[near]
impl Contract {
    /// Read a single value from data storage
    pub fn get_data(&self, key: String) -> String {
        self.kernel_state
            .data_storage
            .get(&key)
            .map(|s| s.clone())
            .unwrap_or_else(|| String::new())
    }

    /// Read multiple values from data storage
    pub fn get_batch_data(&self, keys: Vec<String>) -> Vec<String> {
        keys.into_iter()
            .map(|key| {
                self.kernel_state
                    .data_storage
                    .get(&key)
                    .map(|s| s.clone())
                    .unwrap_or_else(|| String::new())
            })
            .collect()
    }

    /// Get the current lock holder
    pub fn get_lock_holder(&self) -> Option<AccountId> {
        self.kernel_state.global_lock_holder.clone()
    }

    /// Get all pending actions
    pub fn get_pending_actions(&self) -> Vec<String> {
        self.kernel_state
            .pending_actions
            .iter()
            .map(|s| s.clone())
            .collect()
    }

    /// Get vault creation timestamp
    pub fn get_created_at(&self) -> U128 {
        self.created_at
    }

    /// Get time elapsed since vault creation
    pub fn get_time_since_created(&self) -> U128 {
        U128(self.internal_get_time_since_created())
    }

    pub fn get_time_since_last_executed_proposal(&self) -> U128 {
        let current_time = near_sdk::env::block_timestamp() as u128;
        let last_executed = self.kernel_state.last_executed_proposal_timestamp.0;
        U128(current_time.checked_sub(last_executed).unwrap_or(0))
    }

    pub fn get_policy_change_control(&self) -> ChangeControl {
        self.policy_change_control.clone()
    }

    #[private]
    #[init(ignore_state)]
    pub fn migrate() -> Self {
        let mut old_contract: Contract = env::state_read().expect("read state failed");

        old_contract.staged_wasm.set(None);
        old_contract.staged_code_hash = None;

        old_contract
    }
}