use crate::dna::mds::codegen::{HelixIR, Instruction};
use std::collections::{HashMap, HashSet};
pub use crate::dna::mds::codegen::{StringPool, SymbolTable, Metadata, ConstantPool, ConstantValue};
use std::path::PathBuf;
use anyhow::Result;


#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OptimizationLevel {
    Zero,
    One,
    Two,
    Three,
}
impl Default for OptimizationLevel {
    fn default() -> Self {
        Self::Two
    }
}
impl From<u8> for OptimizationLevel {
    fn from(level: u8) -> Self {
        match level {
            0 => Self::Zero,
            1 => Self::One,
            2 => Self::Two,
            3 | _ => Self::Three,
        }
    }
}
pub struct Optimizer {
    level: OptimizationLevel,
    stats: OptimizationStats,
}
impl Optimizer {
    pub fn new(level: OptimizationLevel) -> Self {
        Self {
            level,
            stats: OptimizationStats::default(),
        }
    }
    pub fn optimize(&mut self, ir: &mut HelixIR) {
        match self.level {
            OptimizationLevel::Zero => {}
            OptimizationLevel::One => {
                self.apply_basic_optimizations(ir);
            }
            OptimizationLevel::Two => {
                self.apply_basic_optimizations(ir);
                self.apply_standard_optimizations(ir);
            }
            OptimizationLevel::Three => {
                self.apply_basic_optimizations(ir);
                self.apply_standard_optimizations(ir);
                self.apply_aggressive_optimizations(ir);
            }
        }
    }
    fn apply_basic_optimizations(&mut self, ir: &mut HelixIR) {
        self.deduplicate_strings(ir);
        self.remove_dead_code(ir);
        self.optimize_string_pool(ir);
    }
    fn apply_standard_optimizations(&mut self, ir: &mut HelixIR) {
        self.fold_constants(ir);
        self.inline_small_functions(ir);
        self.optimize_instruction_sequence(ir);
        self.merge_duplicate_sections(ir);
    }
    fn apply_aggressive_optimizations(&mut self, ir: &mut HelixIR) {
        self.eliminate_cross_references(ir);
        self.optimize_pipelines(ir);
        self.compress_data_sections(ir);
        self.reorder_for_cache_locality(ir);
    }
    fn deduplicate_strings(&mut self, ir: &mut HelixIR) {
        let mut seen = HashMap::new();
        let mut new_strings = Vec::new();
        let mut remap = HashMap::new();
        for (idx, string) in ir.string_pool.strings.iter().enumerate() {
            if let Some(&existing_idx) = seen.get(string) {
                remap.insert(idx as u32, existing_idx);
                self.stats.strings_deduplicated += 1;
            } else {
                let new_idx = new_strings.len() as u32;
                seen.insert(string.clone(), new_idx);
                new_strings.push(string.clone());
                remap.insert(idx as u32, new_idx);
            }
        }
        let original_size = ir.string_pool.strings.len();
        ir.string_pool.strings = new_strings;
        self.stats.strings_removed = original_size - ir.string_pool.strings.len();
        for instruction in &mut ir.instructions {
            self.remap_instruction_strings(instruction, &remap);
        }
    }
    fn remove_dead_code(&mut self, ir: &mut HelixIR) {
        let mut reachable = HashSet::new();
        let mut work_list = vec![0];
        while let Some(idx) = work_list.pop() {
            if idx >= ir.instructions.len() || !reachable.insert(idx) {
                continue;
            }
            work_list.push(idx + 1);
        }
        let mut new_instructions = Vec::new();
        let mut remap = HashMap::new();
        for (idx, instruction) in ir.instructions.iter().enumerate() {
            if reachable.contains(&idx) {
                remap.insert(idx, new_instructions.len());
                new_instructions.push(instruction.clone());
            } else {
                self.stats.instructions_removed += 1;
            }
        }
        for instruction in &mut new_instructions {
            self.remap_jump_targets(instruction, &remap);
        }
        ir.instructions = new_instructions;
    }
    fn optimize_string_pool(&mut self, ir: &mut HelixIR) {
        let mut frequency = HashMap::new();
        for instruction in &ir.instructions {
            self.count_string_usage(instruction, &mut frequency);
        }
        let mut indexed_strings: Vec<(u32, String)> = ir
            .string_pool
            .strings
            .iter()
            .enumerate()
            .map(|(i, s)| (i as u32, s.clone()))
            .collect();
        indexed_strings
            .sort_by_key(|(idx, _)| {
                std::cmp::Reverse(frequency.get(idx).copied().unwrap_or(0))
            });
        let mut remap = HashMap::new();
        let mut new_strings = Vec::new();
        for (old_idx, string) in indexed_strings {
            let new_idx = new_strings.len() as u32;
            remap.insert(old_idx, new_idx);
            new_strings.push(string);
        }
        ir.string_pool.strings = new_strings;
        for instruction in &mut ir.instructions {
            self.remap_instruction_strings(instruction, &remap);
        }
    }
    fn fold_constants(&mut self, _ir: &mut HelixIR) {}
    fn inline_small_functions(&mut self, ir: &mut HelixIR) {
        let mut reference_count = std::collections::HashMap::new();
        for instruction in &ir.instructions {
            if let Instruction::ResolveReference { index, .. } = instruction {
                *reference_count.entry(*index).or_insert(0) += 1;
            }
        }
        self.stats.functions_inlined = 0;
    }
    fn optimize_instruction_sequence(&mut self, ir: &mut HelixIR) {
        let mut seen_properties: std::collections::HashSet<(u32, u32)> = std::collections::HashSet::new();
        let mut i = 0;
        while i < ir.instructions.len() {
            match &ir.instructions[i] {
                Instruction::SetProperty { target, key, .. } => {
                    let prop_key = (*target, *key);
                    if seen_properties.contains(&prop_key) {
                        ir.instructions.remove(i);
                        self.stats.instructions_removed += 1;
                        continue;
                    } else {
                        seen_properties.insert(prop_key);
                    }
                }
                Instruction::SetCapability { agent, capability } => {
                    let cap_key = (*agent, *capability);
                    if seen_properties.contains(&cap_key) {
                        ir.instructions.remove(i);
                        self.stats.instructions_removed += 1;
                        continue;
                    } else {
                        seen_properties.insert(cap_key);
                    }
                }
                _ => {}
            }
            i += 1;
        }
    }
    fn merge_duplicate_sections(&mut self, ir: &mut HelixIR) {
        use std::collections::HashMap;
        let mut agent_signatures: HashMap<String, Vec<u32>> = HashMap::new();
        for (id, agent) in &ir.symbol_table.agents {
            let signature = format!(
                "{}-{}-{:?}-{:?}", agent.model_idx, agent.role_idx, agent.temperature,
                agent.max_tokens
            );
            agent_signatures.entry(signature).or_insert_with(Vec::new).push(*id);
        }
        for (_, agents) in &agent_signatures {
            if agents.len() > 1 {
                self.stats.sections_merged += agents.len() - 1;
            }
        }
        let mut workflow_signatures: HashMap<String, Vec<u32>> = HashMap::new();
        for (id, workflow) in &ir.symbol_table.workflows {
            let signature = format!("{:?}", workflow.trigger_type);
            workflow_signatures.entry(signature).or_insert_with(Vec::new).push(*id);
        }
    }
    fn eliminate_cross_references(&mut self, ir: &mut HelixIR) {
        use std::collections::HashSet;
        let mut referenced_agents: HashSet<u32> = HashSet::new();
        let mut referenced_workflows: HashSet<u32> = HashSet::new();
        for crew in ir.symbol_table.crews.values() {
            for agent_id in &crew.agent_ids {
                referenced_agents.insert(*agent_id);
            }
        }
        for workflow in ir.symbol_table.workflows.values() {
            if let Some(pipeline) = &workflow.pipeline {
                for node_id in pipeline {
                    referenced_workflows.insert(*node_id);
                }
            }
        }
        for instruction in &ir.instructions {
            match instruction {
                Instruction::ResolveReference { ref_type: _, index: _ } => {}
                _ => {}
            }
        }
        let unreferenced_agents: Vec<u32> = ir
            .symbol_table
            .agents
            .keys()
            .filter(|id| !referenced_agents.contains(id))
            .cloned()
            .collect();
        for agent_id in unreferenced_agents {
            ir.symbol_table.agents.remove(&agent_id);
            self.stats.instructions_removed += 1;
        }
        let unreferenced_workflows: Vec<u32> = ir
            .symbol_table
            .workflows
            .keys()
            .filter(|id| !referenced_workflows.contains(id))
            .cloned()
            .collect();
        for workflow_id in unreferenced_workflows {
            ir.symbol_table.workflows.remove(&workflow_id);
            self.stats.instructions_removed += 1;
        }
    }
    fn optimize_pipelines(&mut self, ir: &mut HelixIR) {
        for i in 0..ir.instructions.len() {
            if let Instruction::DefinePipeline { .. } = &ir.instructions[i] {
                self.stats.pipelines_optimized += 1;
            }
        }
    }
    fn compress_data_sections(&mut self, ir: &mut HelixIR) {
        let mut string_frequency: HashMap<u32, usize> = HashMap::new();
        for instruction in &ir.instructions {
            match instruction {
                Instruction::SetProperty { key, .. } => {
                    *string_frequency.entry(*key).or_insert(0) += 1;
                }
                Instruction::SetCapability { capability, .. } => {
                    *string_frequency.entry(*capability).or_insert(0) += 1;
                }
                Instruction::SetMetadata { key, value } => {
                    *string_frequency.entry(*key).or_insert(0) += 1;
                    *string_frequency.entry(*value).or_insert(0) += 1;
                }
                _ => {}
            }
        }
        for agent in ir.symbol_table.agents.values() {
            *string_frequency.entry(agent.name_idx).or_insert(0) += 1;
            *string_frequency.entry(agent.model_idx).or_insert(0) += 1;
            *string_frequency.entry(agent.role_idx).or_insert(0) += 1;
        }
        let _total_strings = ir.string_pool.strings.len();
        let frequently_used = string_frequency
            .iter()
            .filter(|(_, count)| **count > 1)
            .count();
        if frequently_used > 0 {
            self.stats.bytes_saved += frequently_used * 8;
        }
    }
    fn reorder_for_cache_locality(&mut self, ir: &mut HelixIR) {
        let mut reordered = Vec::new();
        let mut agent_instructions = Vec::new();
        let mut workflow_instructions = Vec::new();
        let mut other_instructions = Vec::new();
        for instruction in ir.instructions.drain(..) {
            match &instruction {
                Instruction::DeclareAgent(_) => agent_instructions.push(instruction),
                Instruction::DeclareWorkflow(_) => {
                    workflow_instructions.push(instruction)
                }
                Instruction::DefinePipeline { .. } => {
                    workflow_instructions.push(instruction)
                }
                _ => other_instructions.push(instruction),
            }
        }
        reordered.extend(agent_instructions);
        reordered.extend(workflow_instructions);
        reordered.extend(other_instructions);
        ir.instructions = reordered;
    }
    fn remap_instruction_strings(
        &self,
        instruction: &mut Instruction,
        remap: &HashMap<u32, u32>,
    ) {
        match instruction {
            Instruction::SetProperty { key, value, .. } => {
                if let Some(&new_idx) = remap.get(key) {
                    *key = new_idx;
                }
                if let ConstantValue::String(idx) = value {
                    if let Some(&new_idx) = remap.get(idx) {
                        *idx = new_idx;
                    }
                }
            }
            Instruction::SetCapability { capability, .. } => {
                if let Some(&new_idx) = remap.get(capability) {
                    *capability = new_idx;
                }
            }
            Instruction::SetMetadata { key, value } => {
                if let Some(&new_idx) = remap.get(key) {
                    *key = new_idx;
                }
                if let Some(&new_idx) = remap.get(value) {
                    *value = new_idx;
                }
            }
            _ => {}
        }
    }
    fn remap_jump_targets(
        &self,
        _instruction: &mut Instruction,
        _remap: &HashMap<usize, usize>,
    ) {}
    fn count_string_usage(
        &self,
        instruction: &Instruction,
        frequency: &mut HashMap<u32, usize>,
    ) {
        match instruction {
            Instruction::SetProperty { key, value, .. } => {
                *frequency.entry(*key).or_insert(0) += 1;
                if let ConstantValue::String(idx) = value {
                    *frequency.entry(*idx).or_insert(0) += 1;
                }
            }
            Instruction::SetCapability { capability, .. } => {
                *frequency.entry(*capability).or_insert(0) += 1;
            }
            Instruction::SetMetadata { key, value } => {
                *frequency.entry(*key).or_insert(0) += 1;
                *frequency.entry(*value).or_insert(0) += 1;
            }
            _ => {}
        }
    }
    pub fn stats(&self) -> &OptimizationStats {
        &self.stats
    }
}
#[derive(Debug, Default)]
pub struct OptimizationStats {
    pub strings_deduplicated: usize,
    pub strings_removed: usize,
    pub instructions_removed: usize,
    pub constants_folded: usize,
    pub functions_inlined: usize,
    pub pipelines_optimized: usize,
    pub sections_merged: usize,
    pub bytes_saved: usize,
}
impl OptimizationStats {
    pub fn report(&self) -> String {
        format!(
            "Optimization Results:\n\
             - Strings deduplicated: {}\n\
             - Strings removed: {}\n\
             - Instructions removed: {}\n\
             - Constants folded: {}\n\
             - Functions inlined: {}\n\
             - Pipelines optimized: {}\n\
             - Sections merged: {}\n\
             - Total bytes saved: {}",
            self.strings_deduplicated, self.strings_removed, self.instructions_removed,
            self.constants_folded, self.functions_inlined, self.pipelines_optimized, self
            .sections_merged, self.bytes_saved
        )
    }
}
#[cfg(test)]
mod tests {
    use crate::mds::codegen::{StringPool, SymbolTable, Metadata, ConstantPool, ConstantValue};
    #[test]
    fn test_optimization_levels() {
        assert_eq!(OptimizationLevel::from(0), OptimizationLevel::Zero);
        assert_eq!(OptimizationLevel::from(1), OptimizationLevel::One);
        assert_eq!(OptimizationLevel::from(2), OptimizationLevel::Two);
        assert_eq!(OptimizationLevel::from(3), OptimizationLevel::Three);
        assert_eq!(OptimizationLevel::from(99), OptimizationLevel::Three);
    }
    #[test]
    fn test_string_deduplication() {
        let mut ir = HelixIR {
            version: 1,
            metadata: Metadata::default(),
            symbol_table: SymbolTable::default(),
            instructions: vec![
                Instruction::DeclareAgent(0), Instruction::DeclareWorkflow(1),
                Instruction::DeclareContext(2),
            ],
            string_pool: StringPool {
                strings: vec![
                    "hello".to_string(), "world".to_string(), "hello".to_string(),
                ],
                index: std::collections::HashMap::new(),
            },
            constants: ConstantPool::default(),
        };
        let mut optimizer = Optimizer::new(OptimizationLevel::One);
        optimizer.deduplicate_strings(&mut ir);
        assert_eq!(ir.string_pool.strings.len(), 2);
        assert_eq!(optimizer.stats.strings_deduplicated, 1);
    }
    #[test]
    fn test_constant_folding() {
        let mut ir = crate::dna::atp::types::HelixIR {
            version: 1,
            metadata: Metadata::default(),
            symbol_table: SymbolTable::default(),
            instructions: vec![Instruction::DeclareAgent(0),],
            string_pool: StringPool::default(),
            constants: ConstantPool::default(),
        };
        let mut optimizer = Optimizer::new(OptimizationLevel::Two);
        optimizer.fold_constants(&mut ir);
        assert_eq!(ir.instructions.len(), 1);
        match &ir.instructions[0] {
            Instruction::DeclareAgent(0) => {}
            _ => panic!("Expected DeclareAgent(0)"),
        }
    }
}

pub fn optimize_command(
    input: PathBuf,
    output: Option<PathBuf>,
    level: u8,
    verbose: bool,
) -> Result<(), Box<dyn std::error::Error>> {
    
    let output_path = output.unwrap_or_else(|| input.clone());
    if verbose {
        println!("⚡ Optimizing: {}", input.display());
        println!("  Level: {}", level);
    }
    let loader = crate::dna::mds::loader::BinaryLoader::new();
    let binary = loader.load_file(&input)?;
    let serializer = crate::dna::mds::serializer::BinarySerializer::new(false);
    let mut ir = serializer.deserialize_to_ir(&binary)?;
    let mut optimizer = crate::dna::mds::optimizer::Optimizer::new(
        OptimizationLevel::from(level)
    );
    optimizer.optimize(&mut ir);
    let optimized_binary = serializer.serialize(ir, None)?;
    serializer.write_to_file(&optimized_binary, &output_path)?;
    println!("✅ Optimized successfully: {}", output_path.display());
    if verbose {
        let stats = optimizer.stats();
        println!("\nOptimization Results:");
        println!("{}", stats.report());
    }
    Ok(())
}