source_viewer 0.4.3

use crate::file_parser::EStr;
use addr2line::FrameIter;
use once_cell::unsync::OnceCell;
use crate::file_parser::map_dissasm;
use capstone::Capstone;
use std::rc::Rc;
use crate::file_parser::CodeRange;
use std::fmt::Write;
use crate::errors::StackedError;
use crate::errors::WrapedError;
use crate::file_parser::InstructionDetail;
use crate::file_parser::MachineFile;
use addr2line::LookupContinuation;
use addr2line::LookupResult;
use std::collections::{BTreeMap, HashMap};
use std::error::Error;
use std::fs;
use std::path::Path;
use std::sync::Arc;

use std::collections::hash_map;
use typed_arena::Arena;

//probably needed to handle the suplementry matrial

pub struct FileRegistry<'a> {
    pub files_arena: &'a Arena<Vec<u8>>,
    pub map: HashMap<Arc<Path>, Result<MachineFile<'a>, WrapedError>>,
}

impl<'a> FileRegistry<'a> {
    pub fn new(files_arena: &'a Arena<Vec<u8>>) -> Self {
        FileRegistry {
            files_arena,
            map: HashMap::new(),
        }
    }

    pub fn get_machine(&mut self, path: Arc<Path>) -> Result<&mut MachineFile<'a>, Box<dyn Error>> {
        //code looks so ugly because we cant pull into a side function or the borrow checker will freak out
        // println!("geting data for {}",path.to_string_lossy());

        match self.map.entry(path.clone()) {
            hash_map::Entry::Occupied(entry) => {
                entry.into_mut().as_mut().map_err(|e| e.clone().into())
            }
            hash_map::Entry::Vacant(entry) => {
                let buffer = match fs::read(&*path) {
                    Ok(x) => x,
                    Err(e) => {
                        return entry
                            .insert(Err(WrapedError::new(Box::new(e))))
                            .as_mut()
                            .map_err(|e| e.clone().into())
                    }
                };
                let b = self.files_arena.alloc(buffer);
                entry
                    .insert(MachineFile::parse(b).map_err(WrapedError::new))
                    .as_mut()
                    .map_err(|e| e.clone().into())
            }
        }
    }
}

// pub type AddressFileMapping = HashMap<u64, (String, u32)>; // address -> (file, line)

// pub fn map_instructions_to_source(
//     machine_file: &MachineFile,
// ) -> Result<AddressFileMapping, Box<dyn Error>> {
//     let mut mapping = AddressFileMapping::new();

//     // Create addr2line context from DWARF data
//     let ctx = machine_file.get_addr2line()?;
//     let arch = machine_file.obj.architecture();

//     // Iterate through each code section and map addresses to source
//     for section in &machine_file.sections {
//         if let Section::Code(code_section) = section {
//             for instruction in code_section.get_asm(arch)?.iter() {
//                 if let Ok(Some(loc)) = ctx.find_location(instruction.address) {
//                     let file = loc.file.unwrap_or("<unknown>").to_string();
//                     let line = loc.line.unwrap_or(0);
//                     mapping.insert(instruction.address, (file, line));
//                 }
//             }
//         }
//     }

//     Ok(mapping)
// }

pub type DebugContext<'a> = addr2line::Context<EStr<'a>>;

fn select_one_func<'a>(
    mut frames: FrameIter<EStr<'a>>,
) -> Option<String> {
    while let Ok(Some(frame)) = frames.next() {
        if let Some(raw) = frame.function {
            if let Ok(name) = raw.demangle() {
                //inner most function is probably the most intresting
                return Some(name.to_string())
            }
        }
    }

    None
}

fn map_frame_func<'a,E>(
    mut frames: FrameIter<EStr<'a>>,
    mut map:impl FnMut(&str)->Result<(),E>,
) -> Result<(),E>{

    while let Ok(Some(frame)) = frames.next() {
        if let Some(raw) = frame.function {
            if let Ok(name) = raw.demangle() {
                map(&name)?
            }
        }
    }

    Ok(())
}

pub fn resolve_func_name(addr2line: &DebugContext, address: u64) -> Option<String> {
    // Start the frame lookup process
    let lookup_result = addr2line.find_frames(address);

    let frames = lookup_result.skip_all_loads().ok()?;
    select_one_func(frames)
}

fn get_func_frames<'a, 'b: 'a,'c>(
    addr2line: &'c DebugContext<'a>,
    registry: &mut FileRegistry<'b>,
    address: u64,
) -> Option<FrameIter<'c,EStr<'a>>>{
    let mut lookup_result = addr2line.find_frames(address);

    loop {
        match lookup_result {
            LookupResult::Load { load, continuation } => {
                // println!("load case {:?} {:?}",load.parent,load.path);

                // Construct the full path for the DWO file if possible
                let dwo_path = load
                    .comp_dir
                    .as_ref()
                    .map(|comp_dir : &EStr | {
                        std::path::PathBuf::from(comp_dir.to_string_lossy().to_string())
                    })
                    .and_then(|comp_dir_path  | {
                        load.path.as_ref().map(|path:&EStr| {
                            comp_dir_path
                                .join(std::path::Path::new(&path.to_string_lossy().to_string()))
                        })
                    });

                // println!("load case {:?}",dwo_path);

                let dwo = dwo_path.and_then(
                    |full_path:std::path::PathBuf| {
                        registry
                            .get_machine(full_path.into())
                            .ok()
                            .and_then(|m| m.load_dwarf().ok())
                    }, // .map(Arc::new)
                );

                // Resume the lookup with the loaded data
                lookup_result = continuation.resume(dwo);
            }
            LookupResult::Output(Ok(frames)) => {
                // println!("existing case");

                return Some(frames);
            }
            LookupResult::Output(Err(_e)) => {
                // println!("error case {}",e);

                return None;
            }
        }
    }
}

pub fn find_func_name<'a, 'b: 'a>(
    addr2line: &DebugContext<'a>,
    registry: &mut FileRegistry<'b>,
    address: u64,
) -> Option<String> {
    get_func_frames(addr2line,registry,address)
    .map(select_one_func)?
}

pub fn map_funcs<'a, 'b: 'a, E>(
    addr2line: &DebugContext<'a>,
    registry: &mut FileRegistry<'b>,
    address: u64,
    map:impl FnMut(&str)->Result<(),E>,
) -> Result<(),E> {
    if let Some(frame) = get_func_frames(addr2line,registry,address){
        map_frame_func(frame,map)
    }else{
        Ok(())
    }
}

// #[derive(PartialEq,Clone)]
// pub struct Instruction{
//     pub detail:InstructionDetail,
//     pub file: Arc<Path>
// }

pub struct LazeyAsm<'a>{
    ranges: Vec<CodeRange<'a>>,
    cs:Rc<Capstone>,
    asm:OnceCell<Box<[InstructionDetail]>>
}

impl<'a> LazeyAsm<'a>{
    pub fn new(cs:Rc<Capstone>)->Self{
        Self{
            ranges:Vec::new(),
            asm:OnceCell::new(),
            cs,
        }
    }

    pub fn make_asm(&self)->Result<&[InstructionDetail],Box<dyn Error>>{
        self.asm.get_or_try_init(||{
            let mut ans = Vec::new();
            for r in &self.ranges{
                map_dissasm(
                    &self.cs,
                    r.data,
                    r.address,
                    &mut |ins|{Ok(ans.push(ins))}
                )?;
            }
            Ok(ans.into())
        }).map(|b|&**b)
        
    }
}

// #[derive(PartialEq)]
pub struct CodeFile<'a> {
    pub text: String,
    line_map:OnceCell<HashMap<u32,(usize,usize)>>,//line->byte span
    asm: BTreeMap<u32, HashMap<Arc<Path>, LazeyAsm<'a>>>, //line -> instruction
    pub errors: Vec<(StackedError, Option<Arc<Path>>)>,
}

impl<'a> CodeFile<'a> {
    pub fn read(path: &Path) -> Result<Self, Box<dyn Error>> {
        let text = fs::read_to_string(path)?;
        Ok(CodeFile {
            text,
            line_map:OnceCell::new(),
            asm: BTreeMap::new(),
            errors: Vec::new(),
        })
    }

    pub fn read_arena<'r>(
        path: &Path,
        arena: &'r Arena<CodeFile<'a>>,
    ) -> Result<&'r mut Self, Box<dyn Error>> {
        Ok(arena.alloc(CodeFile::read(path)?))
    }

    // #[inline]
    // pub fn get_asm(&self, line: &u32, obj_path: Arc<Path>) -> Option<&[InstructionDetail]> {
    //     self.asm.get(line)?.get(&obj_path).map(|x| x.as_slice()) //.unwrap_or(&[])
    // }

    #[inline]
    pub fn get_asm(&self, line: &u32, obj_path: Arc<Path>) -> Option<Result<&[InstructionDetail],Box<dyn Error>>> {
        self.asm.get(line)?.get(&obj_path).map(|x| x.make_asm()) //.unwrap_or(&[])
    }

    #[inline]
    pub fn get_line(&self,line:u32)->Option<&str>{
        let (start,end) = self.get_line_map().get(&line)?;
        let slice = &self.text.as_bytes()[*start..*end];
        Some(std::str::from_utf8(slice).unwrap())
    }

    fn get_line_map(&self)->&HashMap<u32,(usize,usize)>{
        self.line_map.get_or_init(||{
            self.text.lines().enumerate().map(|(i,t)|{
                let number = i as u32 + 1;
                let text_start = t.as_ptr().addr()-self.text.as_ptr().addr();
                let text_end = text_start+t.len();
                (number,(text_start,text_end))
            }).collect()
        })
    }

    fn populate(&mut self, asm: &mut FileRegistry<'a>, path: Arc<Path>){
        // Helper closure that runs a fallible block, catches any Err,
        // pushes to self.errors, and continues the outer loop.
        macro_rules! try_wrapped {
            ($expr:expr, $msg:expr) => {
                match $expr {
                    Ok(v) => v,
                    Err(e) => {
                        let err = StackedError::new(e, $msg);
                        self.errors.push((err, None));
                        continue;
                    }
                }
            };
        }

        for (obj_path, res) in asm.map.iter_mut() {
            // both can use normal `?` style thanks to the macro
            let machine_file = try_wrapped!(res.as_ref().map_err(|e| e.clone().into()), "while getting machine");
            let cs = try_wrapped!(machine_file.get_capstone(), "while making dissasmbler");
            let map = try_wrapped!(machine_file.get_lines_map(), "while making context");

            if let Some(line_map) = map.get(&path) {
                for (line, v) in line_map.iter_maped() {
                    self
                        .asm
                        .entry(*line)
                        .or_insert_with(HashMap::new)
                        .entry(obj_path.clone())
                        .or_insert_with(||LazeyAsm::new(cs.clone()))
                        .ranges.extend_from_slice(v);
                }
            }
        }
    }

    pub fn get_error(&self)->Result<(),String>{
        if !self.errors.is_empty() {
            let mut output = String::new();

            writeln!(
                &mut output,
                "⚠️  Warning: errors occurred while reading debug info ({} total):",
                self.errors.len()
            ).ok();

            for (err, path_opt) in &self.errors {
                let path_str = path_opt
                    .as_ref()
                    .map(|p| p.display().to_string())
                    .unwrap_or_else(|| "<unknown>".to_string());

                writeln!(&mut output, "• Path: {path_str}\n  Error: {err}\n").ok();
            }

            return Err(output)
        }
        Ok(())
    }

}

pub struct CodeRegistry<'data, 'r> {
    pub source_files: HashMap<Arc<Path>, Result<&'r CodeFile<'data>, Box<WrapedError>>>,
    pub asm: &'r mut FileRegistry<'data>,
    arena: &'r Arena<CodeFile<'data>>,
    // pub visited : HashSet<Arc<Path>>,
    // pub asm: FileRegistry<'a>,
}

impl<'data, 'r> CodeRegistry<'data, 'r> {
    pub fn new(asm: &'r mut FileRegistry<'data>, arena: &'r Arena<CodeFile<'data>>) -> Self {
        CodeRegistry {
            asm,
            arena,
            source_files: HashMap::new(),
        }
    }

    // pub fn format_inst_debug(&mut self,ins:&InstructionDetail,debug:&DebugContext<'data>)->String{
    //     format_inst_debug(ins,debug,self.asm)
    // }

    pub fn get_existing_source_file(
        &self,
        path: &Arc<Path>,
    ) -> Result<&'r CodeFile<'data>, Box<dyn Error>> {
        self.source_files
            .get(path)
            .unwrap()
            .as_ref()
            .map_err(|e| e.clone().into())
            .copied()
    }

    pub fn get_source_file(&mut self, path: Arc<Path>,dwarf_errors:bool) -> Result<&'r CodeFile<'data>, Box<dyn Error>> {
        match self.source_files.entry(path.clone()) {
            hash_map::Entry::Occupied(entry) => entry.get().clone().map_err(|e| e.clone().into()),
            hash_map::Entry::Vacant(entry) => {
                // let code_file = entry.insert(CodeFile::read_arena(&path,self.arena)
                //     .map_err(|e| Box::new(WrapedError::new(e)))
                // );

                let code_file = match CodeFile::read_arena(&path, self.arena) {
                    Ok(x) => x,
                    Err(e) => {
                        let err = Box::new(WrapedError::new(e));
                        entry.insert(Err(err.clone()));
                        return Err(err);
                    }
                };

                code_file.populate(self.asm, path);
                if dwarf_errors{
                    code_file.get_error()?
                }
                


                entry.insert(Ok(code_file));
                Ok(code_file)
            }
        }
    }

    pub fn visit_machine_file(
        &mut self,
        path: Arc<Path>,
    ) -> Result<&mut MachineFile<'data>, Box<dyn Error>> {
        self.asm.get_machine(path)
    }

    pub fn get_existing_machine(&self,path:&Path)->Option<&MachineFile<'data>>{
        self.asm.map.get(path).map(|x| x.as_ref().ok())?
    }
}

// pub fn format_inst_debug<'a, 'b: 'a, 'c>(
//     ins: &InstructionDetail,
//     addr2line: &'c DebugContext<'a>,
//     registry: &mut FileRegistry<'b>,
// ) -> String {
//     format!(
//         "{:#010x}: {:<6} {:<30} {}",
//         ins.address,
//         ins.mnemonic,
//         ins.op_str, //this needs a fixup
//         find_func_name(addr2line, registry, ins.address).unwrap_or("<unknown>".to_string()),
//     )
// }

// pub struct DebugInstruction<'b, 'a> {
//     ins: InstructionDetail,
//     addr2line: &'b DebugContext<'a>,
//     //needs a way to load the Sup files which are machine files...
//     //probably means we need the asm registry
// }

// impl<'a, 'c> DebugInstruction<'c, 'a> {
//     pub fn new(
//         ins: InstructionDetail,
//         addr2line: &'c DebugContext<'a>,
//     ) -> Self {
//         DebugInstruction { ins, addr2line }
//     }

//     // pub fn get_func_name(&self ) ->Option<String> {
//     //     self.resolve_function_name(self.ins.address)
//     // }

//     pub fn get_string_load<'b: 'a>(&self, registry: &mut FileRegistry<'b>) -> String {
//         format!(
//             "{:#010x}: {:<6} {:<30} {}",
//             self.ins.address,
//             self.ins.mnemonic,
//             self.ins.op_str, //this needs a fixup
//             find_func_name(self.addr2line, registry, self.ins.address)
//                 .unwrap_or("<unknown>".to_string()),
//         )
//     }

//     // pub fn get_string_no_load(&self) -> String {
//     //     format!("{:#010x}: {:<6} {:<30} {}",

//     //         self.ins.address,
//     //         self.ins.mnemonic,
//     //         self.ins.op_str, //this needs a fixup
//     //         self.get_func_name().unwrap_or("<unknown>".to_string()),
//     //     )
//     // }

//     // Resolve the function name for a given address using addr2line

//     // fn resolve_function_name(&self, address: u64) -> Option<String> {
//     //     resolve_func_name(self.addr2line,address)
//     // }
// }