// CARRY Quantum Entangled Multi-Agent Module
// Manages complementary agent pairs (Primary COMPUTE + Partner VERIFY/MONITOR)
// sharing explicitly modeled quantum state & entanglement registers.

use std::collections::HashSet;
use crate::quantum::core::QuantumSimulator;

/// Agent Operational Role.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AgentRole {
    PrimaryCompute,
    PartnerVerify,
    PartnerMonitor,
    PartnerCrossCheck,
    PartnerRecover,
}

/// Agent Audit Event.
#[derive(Debug, Clone)]
pub struct AgentEvent {
    pub cycle: usize,
    pub agent_id: String,
    pub role: AgentRole,
    pub action: String,
    pub target_qubits: Vec<usize>,
    pub timestamp_ms: u128,
    pub status: String,
}

/// Agent definition with formal permissions, ownership, and cycle budget.
#[derive(Debug, Clone)]
pub struct Agent {
    pub id: String,
    pub role: AgentRole,
    pub qubit_ownership: HashSet<usize>,
    pub permissions: HashSet<String>,
    pub cycle_budget: usize,
    pub cycles_used: usize,
    pub partner_id: Option<String>,
}

impl Agent {
    pub fn new_primary(id: impl Into<String>, qubits: Vec<usize>, cycle_budget: usize) -> Self {
        let mut permissions = HashSet::new();
        permissions.insert("STATE_PREPARE".to_string());
        permissions.insert("GATE_EXECUTE".to_string());
        permissions.insert("MEASURE".to_string());

        Self {
            id: id.into(),
            role: AgentRole::PrimaryCompute,
            qubit_ownership: qubits.into_iter().collect(),
            permissions,
            cycle_budget,
            cycles_used: 0,
            partner_id: None,
        }
    }

    pub fn new_partner(id: impl Into<String>, partner_primary_id: impl Into<String>, role: AgentRole, qubits: Vec<usize>) -> Self {
        let mut permissions = HashSet::new();
        permissions.insert("INVARIANT_VERIFY".to_string());
        permissions.insert("CROSS_CHECK".to_string());
        permissions.insert("AUDIT_INSPECT".to_string());
        if role == AgentRole::PartnerRecover {
            permissions.insert("RECOVER_STATE".to_string());
        }

        Self {
            id: id.into(),
            role,
            qubit_ownership: qubits.into_iter().collect(),
            permissions,
            cycle_budget: 1000,
            cycles_used: 0,
            partner_id: Some(partner_primary_id.into()),
        }
    }

    pub fn can_perform(&self, action: &str) -> bool {
        self.permissions.contains(action)
    }

    pub fn owns_qubit(&self, qubit: usize) -> bool {
        self.qubit_ownership.contains(&qubit)
    }

    pub fn consume_cycle(&mut self) -> Result<(), String> {
        if self.cycles_used >= self.cycle_budget {
            return Err(format!("AGENT_BUDGET_EXCEEDED: Agent {} exceeded budget {}", self.id, self.cycle_budget));
        }
        self.cycles_used += 1;
        Ok(())
    }
}

/// Complementary Agent Pair containing Primary (COMPUTE) and Partner (VERIFY/MONITOR).
#[derive(Debug, Clone)]
pub struct AgentPair {
    pub primary: Agent,
    pub partner: Agent,
    pub entanglement_group: String,
}

impl AgentPair {
    pub fn new(primary: Agent, partner: Agent, group: impl Into<String>) -> Result<Self, String> {
        if primary.role != AgentRole::PrimaryCompute {
            return Err("Primary agent must have role PrimaryCompute".into());
        }
        if partner.role == AgentRole::PrimaryCompute {
            return Err("Partner agent cannot have role PrimaryCompute".into());
        }

        Ok(Self {
            primary,
            partner,
            entanglement_group: group.into(),
        })
    }

    /// Perform a cross-verification between Primary computation and Partner monitor.
    pub fn cross_verify_state(&self, sim: &QuantumSimulator) -> Result<String, String> {
        // Partner checks state purity & trace
        let purity = sim.state.purity();
        if (purity - 1.0).abs() > 0.05 && sim.state.is_valid_state() {
            // Mixed state / noise detected by partner
            Ok(format!("PARTNER_VERIFIED: State is mixed (purity {:.4})", purity))
        } else if sim.state.is_valid_state() {
            Ok(format!("PARTNER_VERIFIED: State is pure (purity {:.4})", purity))
        } else {
            Err("PARTNER_REJECT: Invalid quantum state detected during cross-verification".into())
        }
    }
}

/// Multi-Agent Coordinator managing active agents, pairs, and audit records.
#[derive(Debug, Default)]
pub struct MultiAgentCoordinator {
    pub agents: Vec<Agent>,
    pub pairs: Vec<AgentPair>,
    pub audit_log: Vec<AgentEvent>,
}

impl MultiAgentCoordinator {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn register_pair(&mut self, pair: AgentPair) {
        self.agents.push(pair.primary.clone());
        self.agents.push(pair.partner.clone());
        self.pairs.push(pair);
    }

    pub fn log_event(&mut self, event: AgentEvent) {
        self.audit_log.push(event);
    }

    pub fn verify_action_permissions(&self, agent_id: &str, action: &str, qubit: usize) -> Result<(), String> {
        let agent = self
            .agents
            .iter()
            .find(|a| a.id == agent_id)
            .ok_or_else(|| format!("Agent {} not found", agent_id))?;

        if !agent.can_perform(action) {
            return Err(format!("UNAUTHORIZED_ACTION: Agent {} lacks permission for {}", agent_id, action));
        }

        if !agent.owns_qubit(qubit) {
            return Err(format!("UNAUTHORIZED_QUBIT: Agent {} does not own qubit {}", agent_id, qubit));
        }

        Ok(())
    }
}