//! Proposal generation and validation for autonomous optimization.

use crate::error::AutoResult;
use rtx_runtime::Runtime;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use uuid::Uuid;

/// Types of optimization proposals that can be generated.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ProposalType {
    /// Data layout optimization proposals
    DataLayout,
    /// Memory access coalescing proposals
    MemoryCoalescing,
    /// Cache optimization proposals
    CacheOptimization,
    /// Data parallelization proposals
    DataParallel,
    /// Pipeline parallelization proposals
    PipelineParallel,
    /// Tensor parallelization proposals
    TensorParallel,
    /// Communication optimization proposals
    CommunicationOptimization,
    /// Quantization proposals
    Quantization,
    /// Mixed precision proposals
    MixedPrecision,
    /// Calibration proposals
    Calibration,
    /// Adaptive quantization proposals
    AdaptiveQuantization,
    /// Kernel fusion proposals
    KernelFusion,
    /// Kernel optimization proposals
    KernelOptimization,
    /// Memory access pattern optimization
    MemoryAccess,
    /// GPU occupancy optimization
    Occupancy,
    /// Code generation optimization
    CodeGeneration,
}

/// Status of an optimization proposal.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ProposalStatus {
    /// Proposal has been created but not yet evaluated
    Pending,
    /// Proposal is currently being evaluated
    Evaluating,
    /// Proposal has been evaluated and approved for application
    Approved,
    /// Proposal has been rejected
    Rejected,
    /// Proposal has been applied successfully
    Applied,
    /// Proposal application failed
    Failed,
    /// Proposal has been rolled back
    RolledBack,
}

/// An optimization proposal generated by autonomous agents.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Proposal {
    id: String,
    proposal_type: ProposalType,
    description: String,
    expected_performance_gain: f32,
    status: ProposalStatus,
    created_at: u64,
    metadata: HashMap<String, String>,
}

impl Proposal {
    /// Create a new optimization proposal.
    pub fn new(
        proposal_type: ProposalType,
        description: String,
        expected_performance_gain: f32,
    ) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            proposal_type,
            description,
            expected_performance_gain,
            status: ProposalStatus::Pending,
            created_at: std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_secs(),
            metadata: HashMap::new(),
        }
    }

    /// Get the proposal ID.
    pub fn id(&self) -> &str {
        &self.id
    }

    /// Get the proposal type.
    pub fn proposal_type(&self) -> ProposalType {
        self.proposal_type
    }

    /// Get the proposal description.
    pub fn description(&self) -> &str {
        &self.description
    }

    /// Get the expected performance gain.
    pub fn expected_performance_gain(&self) -> f32 {
        self.expected_performance_gain
    }

    /// Get the current status.
    pub fn status(&self) -> ProposalStatus {
        self.status
    }

    /// Get the creation timestamp.
    pub fn created_at(&self) -> u64 {
        self.created_at
    }

    /// Set the proposal status.
    pub fn set_status(&mut self, status: ProposalStatus) {
        self.status = status;
    }

    /// Add metadata to the proposal.
    pub fn add_metadata<K: Into<String>, V: Into<String>>(&mut self, key: K, value: V) {
        self.metadata.insert(key.into(), value.into());
    }

    /// Get metadata value by key.
    pub fn get_metadata(&self, key: &str) -> Option<&String> {
        self.metadata.get(key)
    }
}

/// Validates and scores optimization proposals.
pub struct ProposalValidator {
    runtime: Arc<Runtime>,
}

impl ProposalValidator {
    /// Create a new proposal validator.
    pub fn new(runtime: Arc<Runtime>) -> AutoResult<Self> {
        Ok(Self { runtime })
    }

    /// Validate that the expected performance gain is reasonable.
    pub async fn validate_performance_gain(&self, proposal: &Proposal) -> AutoResult<bool> {
        // Realistic performance gains depend on proposal type
        let max_reasonable_gain = match proposal.proposal_type {
            ProposalType::DataLayout => 5.0,
            ProposalType::MemoryCoalescing => 3.0,
            ProposalType::CacheOptimization => 2.5,
            ProposalType::DataParallel => 8.0,
            ProposalType::PipelineParallel => 6.0,
            ProposalType::TensorParallel => 4.0,
            ProposalType::CommunicationOptimization => 3.0,
            ProposalType::Quantization => 5.0,
            ProposalType::MixedPrecision => 3.5,
            ProposalType::Calibration => 2.0,
            ProposalType::AdaptiveQuantization => 4.0,
            ProposalType::KernelFusion => 8.0,
            ProposalType::KernelOptimization => 10.0,
            ProposalType::MemoryAccess => 4.0,
            ProposalType::Occupancy => 3.0,
            ProposalType::CodeGeneration => 5.0,
        };

        Ok(proposal.expected_performance_gain > 0.0
            && proposal.expected_performance_gain <= max_reasonable_gain)
    }

    /// Validate the feasibility of applying the proposal.
    pub async fn validate_feasibility(&self, proposal: &Proposal) -> AutoResult<bool> {
        // Basic feasibility checks
        if proposal.description.is_empty() {
            return Ok(false);
        }

        if proposal.expected_performance_gain <= 0.0 {
            return Ok(false);
        }

        // More sophisticated feasibility analysis would go here
        // For now, we consider all proposals with reasonable descriptions feasible
        Ok(true)
    }

    /// Score a proposal based on multiple criteria.
    pub async fn score_proposal(&self, proposal: &Proposal) -> AutoResult<f32> {
        let mut score = 0.0;

        // Performance gain contribution (0.0 - 0.4)
        let normalized_gain = (proposal.expected_performance_gain.min(10.0)) / 10.0;
        score += normalized_gain * 0.4;

        // Proposal type priority (0.0 - 0.3)
        let type_priority = match proposal.proposal_type {
            ProposalType::KernelFusion | ProposalType::KernelOptimization => 0.3,
            ProposalType::DataParallel | ProposalType::PipelineParallel => 0.25,
            ProposalType::Quantization | ProposalType::MixedPrecision => 0.2,
            ProposalType::MemoryCoalescing | ProposalType::CacheOptimization => 0.15,
            _ => 0.1,
        };
        score += type_priority;

        // Description quality (0.0 - 0.2)
        let description_score = if proposal.description.len() > 50 {
            0.2
        } else {
            0.1
        };
        score += description_score;

        // Metadata completeness (0.0 - 0.1)
        let metadata_score = if proposal.metadata.is_empty() {
            0.0
        } else {
            0.1
        };
        score += metadata_score;

        Ok(score.min(1.0))
    }

    /// Rank a list of proposals by their scores.
    pub async fn rank_proposals(&self, proposals: &[Proposal]) -> AutoResult<Vec<(Proposal, f32)>> {
        let mut scored_proposals = Vec::new();

        for proposal in proposals {
            let score = self.score_proposal(proposal).await?;
            scored_proposals.push((proposal.clone(), score));
        }

        // Sort by score in descending order
        scored_proposals.sort_by(|a, b| b.1.total_cmp(&a.1));

        Ok(scored_proposals)
    }

    /// Detect conflicts between proposals.
    pub async fn detect_conflicts(
        &self,
        proposals: &[Proposal],
    ) -> AutoResult<Vec<(usize, usize)>> {
        let mut conflicts = Vec::new();

        for (i, proposal1) in proposals.iter().enumerate() {
            for (j, proposal2) in proposals.iter().enumerate().skip(i + 1) {
                if self.proposals_conflict(proposal1, proposal2) {
                    conflicts.push((i, j));
                }
            }
        }

        Ok(conflicts)
    }

    /// Check if two proposals conflict with each other.
    fn proposals_conflict(&self, proposal1: &Proposal, proposal2: &Proposal) -> bool {
        match (proposal1.proposal_type, proposal2.proposal_type) {
            // Data layout optimizations may conflict
            (ProposalType::DataLayout, ProposalType::DataLayout) => {
                proposal1.description != proposal2.description
            }
            // Memory access patterns may conflict
            (ProposalType::MemoryAccess, ProposalType::MemoryAccess) => true,
            (ProposalType::MemoryCoalescing, ProposalType::MemoryAccess) => true,
            (ProposalType::MemoryAccess, ProposalType::MemoryCoalescing) => true,
            // Parallelization strategies may conflict
            (ProposalType::DataParallel, ProposalType::PipelineParallel) => true,
            (ProposalType::PipelineParallel, ProposalType::DataParallel) => true,
            // Quantization strategies may conflict
            (ProposalType::Quantization, ProposalType::MixedPrecision) => true,
            (ProposalType::MixedPrecision, ProposalType::Quantization) => true,
            _ => false,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_proposal_creation() {
        let proposal = Proposal::new(ProposalType::DataLayout, "Test proposal".to_string(), 2.0);

        assert_eq!(proposal.proposal_type, ProposalType::DataLayout);
        assert_eq!(proposal.expected_performance_gain, 2.0);
        assert_eq!(proposal.status, ProposalStatus::Pending);
    }

    #[test]
    fn test_proposal_status_update() {
        let mut proposal =
            Proposal::new(ProposalType::KernelFusion, "Fuse kernels".to_string(), 3.0);

        proposal.set_status(ProposalStatus::Approved);
        assert_eq!(proposal.status, ProposalStatus::Approved);
    }

    #[test]
    fn test_proposal_metadata() {
        let mut proposal = Proposal::new(
            ProposalType::Quantization,
            "Apply quantization".to_string(),
            2.5,
        );

        proposal.add_metadata("precision", "int8");
        assert_eq!(
            proposal.get_metadata("precision"),
            Some(&"int8".to_string())
        );
    }
}