use colored::Colorize;
use regex::Regex;
use std::collections::{HashMap, HashSet};

use crate::ir;

pub fn run(file: &str, output_ll: bool) -> i32 {
    let content = match std::fs::read_to_string(file) {
        Ok(c) => c,
        Err(e) => {
            eprintln!("{}: Failed to read {}: {}", "error".red().bold(), file, e);
            return 1;
        }
    };

    let module = ir::parse_ir(&content);

    // Collect defined types
    let type_def_re = Regex::new(r"^(%(?:struct|enum)\.\w+)\s*=\s*type\s").unwrap();
    let mut defined: HashSet<String> = HashSet::new();
    for line in content.lines() {
        if let Some(caps) = type_def_re.captures(line) {
            defined.insert(caps[1].to_string());
        }
    }

    // Collect all used named types
    let type_use_re = Regex::new(r"(%(?:struct|enum)\.\w+)").unwrap();
    let mut used: HashSet<String> = HashSet::new();
    for line in content.lines() {
        for caps in type_use_re.captures_iter(line) {
            used.insert(caps[1].to_string());
        }
    }

    let missing: Vec<String> = used.difference(&defined).cloned().collect();
    if missing.is_empty() {
        println!("{} All types defined — nothing to infer.", "culebra".green().bold());
        return 0;
    }

    // Infer types from insertvalue chains in functions
    // Pattern: %name.fN = insertvalue %struct.X %prev, <type> %val, N
    let iv_re = Regex::new(
        r"insertvalue (%(?:struct|enum)\.\w+) (?:%[\w.]+|undef), (.+?), (\d+)"
    ).unwrap();

    // Collect field types per struct: struct_name -> { index -> type }
    let mut field_map: HashMap<String, HashMap<usize, String>> = HashMap::new();

    for func in module.functions.values() {
        for line in func.body.lines() {
            let trimmed = line.trim();
            if let Some(caps) = iv_re.captures(trimmed) {
                let struct_name = caps[1].to_string();
                if !missing.contains(&struct_name) {
                    continue;
                }
                let field_type = clean_field_type(&caps[2]);
                let index: usize = caps[3].parse().unwrap_or(0);
                field_map.entry(struct_name)
                    .or_default()
                    .entry(index)
                    .or_insert(field_type);
            }
        }
    }

    // Also check sret function signatures for return types
    let sret_re = Regex::new(r"sret\((%(?:struct|enum)\.\w+)\)").unwrap();
    // And alloca types
    let alloca_re = Regex::new(r"alloca (%(?:struct|enum)\.\w+)").unwrap();

    let inferred_count = field_map.len();
    let no_info: Vec<_> = missing.iter().filter(|m| !field_map.contains_key(*m)).collect();

    if output_ll {
        // Output as valid LLVM IR type definitions
        println!("; Inferred type definitions for {} (generated by culebra infer-types)", file);
        println!("; {} types inferred from insertvalue chains", inferred_count);
        println!();
    } else {
        println!(
            "{} Inferred {} of {} missing types from insertvalue chains",
            "culebra".green().bold(),
            inferred_count,
            missing.len()
        );
        println!();
    }

    // Sort by name for stable output
    let mut sorted: Vec<_> = field_map.into_iter().collect();
    sorted.sort_by(|a, b| a.0.cmp(&b.0));

    for (struct_name, fields) in &sorted {
        let max_idx = fields.keys().max().copied().unwrap_or(0);
        let mut field_types: Vec<String> = Vec::new();

        for i in 0..=max_idx {
            if let Some(ty) = fields.get(&i) {
                field_types.push(ty.clone());
            } else {
                // Gap — unknown field, use i64 as placeholder
                field_types.push("i64".to_string());
            }
        }

        let type_def = format!("{{ {} }}", field_types.join(", "));

        if output_ll {
            println!("{} = type {}", struct_name, type_def);
        } else {
            println!(
                "  {} = type {} ({} fields)",
                struct_name.yellow().bold(),
                type_def,
                field_types.len()
            );
        }
    }

    // Report types we couldn't infer
    if !no_info.is_empty() {
        if output_ll {
            println!();
            println!("; Could not infer ({} types — no insertvalue chains found):", no_info.len());
            for t in &no_info {
                println!("; {}", t);
            }
        } else {
            println!();
            println!(
                "  {} {} types with no insertvalue chains (can't infer):",
                "⚠".yellow().bold(),
                no_info.len()
            );
            for t in &no_info {
                println!("    {}", t.dimmed());
            }
        }
    }

    if output_ll {
        println!();
        println!("; Usage: prepend these definitions to your .ll file, or add to fix_stage2_phis.py");
    }

    0
}

/// Clean a field type string from insertvalue operand
fn clean_field_type(raw: &str) -> String {
    let raw = raw.trim();
    // Remove the value name: "{ ptr, i64 } %name_val0" -> "{ ptr, i64 }"
    // The type is everything before the last %name or @name
    let re = Regex::new(r"^(.+?)\s+[%@][\w.]+$").unwrap();
    if let Some(caps) = re.captures(raw) {
        return caps[1].trim().to_string();
    }
    // Might be a literal: "i64 0" -> "i64"
    let lit_re = Regex::new(r"^(\S+)\s+\d+$").unwrap();
    if let Some(caps) = lit_re.captures(raw) {
        return caps[1].trim().to_string();
    }
    // Already just a type
    raw.to_string()
}
