use crate::agent::{verification, Agent};
use crate::permission::types::{Codehash, RoleTarget};
use crate::*;
use near_sdk::{env, near, require};

#[near]
impl Contract {
    /// Register an agent with SGX/TDX quote verification
    pub fn register_agent(
        &mut self,
        quote_hex: String,
        collateral: String,
        checksum: String,
        tcb_info: String,
    ) -> bool {
        let (shade_agent_api_image, shade_agent_app_image) =
            verification::verify_agent_quote(quote_hex, collateral, checksum.clone(), tcb_info);

        // verify the shade_agent_api_image codehash has an assigned role (use a role not associated with any policy)
        let api_codehash_target =
            RoleTarget::Codehash(Codehash::from(shade_agent_api_image.clone()));
        require!(
            !self.get_roles_for_target(api_codehash_target).is_empty(),
            "Codehash must have assigned role"
        );
        // verify the shade_agent_app_image codehash has an assigned role (use actual roles shade agent is expected to have)
        let app_codehash_target =
            RoleTarget::Codehash(Codehash::from(shade_agent_app_image.clone()));
        require!(
            !self.get_roles_for_target(app_codehash_target).is_empty(),
            "Codehash must have assigned role"
        );

        let predecessor = env::predecessor_account_id();
        self.kernel_state.registered_agent_by_account_id.insert(
            predecessor,
            Agent {
                checksum,
                codehash: shade_agent_app_image,
            },
        );

        true
    }

    /// Get agent information for an account
    pub fn get_agent(&self, account_id: AccountId) -> Agent {
        self.kernel_state
            .registered_agent_by_account_id
            .get(&account_id)
            .expect("no worker found")
            .to_owned()
    }
}

impl Contract {
    /// Check if account has a codehash with assigned role (internal validation)
    pub fn internal_validate_agent_with_role(&self) {
        let agent = self.get_agent(env::predecessor_account_id());
        let codehash_target = RoleTarget::Codehash(Codehash::from(agent.codehash));
        require!(
            !self.get_roles_for_target(codehash_target).is_empty(),
            "Agent codehash must have assigned role"
        );
    }
}