ghidra 0.0.3

use std::{marker::PhantomData, path::Path};

use crate::{
    Address, AddressRange, AnalysisOptions, AnalysisReport, ControlFlowGraph, DataInfo,
    DecompileOptions, DecompileResult, Error, ExportInfo, FunctionControlFlowGraph,
    FunctionDescriptor, FunctionSummary, GhidraRuntime, ImportInfo, InstructionInfo, LaunchOptions,
    MemoryBlockInfo, ProgramCallGraph, ProgramLoadInfo, ProgramLoadOptions, ProgramMetadata,
    ProgramOpenOptions, ReferenceInfo, Result, SourceMapInfo, SymbolInfo,
    bridge::{self, JavaHandle},
};

/// Entry point for project and program operations.
#[derive(Clone)]
pub struct Ghidra {
    runtime: GhidraRuntime,
}

/// Open Ghidra project handle.
pub struct Project {
    runtime: GhidraRuntime,
    handle: Option<JavaHandle>,
}

/// Open Ghidra program handle tied to its owning project.
pub struct Program<'project> {
    runtime: GhidraRuntime,
    project: JavaHandle,
    handle: Option<JavaHandle>,
    _project: PhantomData<&'project Project>,
}

/// Result of opening or importing a program.
pub struct LoadedProgram<'project> {
    pub program: Program<'project>,
    pub info: ProgramLoadInfo,
}

/// Open Ghidra transaction.
///
/// Dropping an uncommitted transaction rolls it back.
pub struct ProgramTransaction<'program> {
    runtime: GhidraRuntime,
    handle: Option<JavaHandle>,
    _program: PhantomData<&'program ()>,
}

/// Reusable decompiler session for a program.
///
/// Reusing one session preserves the program type registry used by decompile
/// results and `program_metadata`.
pub struct Decompiler<'program> {
    runtime: GhidraRuntime,
    handle: Option<JavaHandle>,
    _program: PhantomData<&'program ()>,
}

impl Ghidra {
    /// Starts the embedded Ghidra runtime and attaches the bridge classes.
    pub fn start(options: LaunchOptions) -> Result<Self> {
        Ok(Self {
            runtime: bridge::start(options)?,
        })
    }

    /// Builds a Ghidra API facade from an already-started runtime.
    pub fn with_runtime(runtime: GhidraRuntime) -> Result<Self> {
        bridge::attach_bridge(&runtime)?;
        Ok(Self { runtime })
    }

    /// Returns the embedded runtime used by this facade.
    pub fn runtime(&self) -> &GhidraRuntime {
        &self.runtime
    }

    /// Opens an existing project or creates it at `project_location`.
    pub fn open_or_create_project(
        &self,
        project_location: impl AsRef<Path>,
        project_name: &str,
    ) -> Result<Project> {
        let handle =
            bridge::open_or_create_project(&self.runtime, project_location.as_ref(), project_name)?;
        Ok(Project {
            runtime: self.runtime.clone(),
            handle: Some(handle),
        })
    }
}

impl Project {
    /// Opens an existing root-level program matching the binary file name or imports the binary.
    pub fn open_or_import_program(&self, binary: impl AsRef<Path>) -> Result<Program<'_>> {
        Ok(self
            .open_or_import_program_with_options(ProgramLoadOptions::new(binary.as_ref())?)?
            .program)
    }

    /// Opens a program already present in this project.
    pub fn open_program(&self, options: ProgramOpenOptions) -> Result<Program<'_>> {
        let project = self.handle()?;
        let handle = bridge::open_program(&self.runtime, project, &options.path)?;
        Ok(Program {
            runtime: self.runtime.clone(),
            project,
            handle: Some(handle),
            _project: PhantomData,
        })
    }

    /// Opens or imports a program with explicit loader, path, and monitor options.
    pub fn open_or_import_program_with_options(
        &self,
        options: ProgramLoadOptions,
    ) -> Result<LoadedProgram<'_>> {
        let project = self.handle()?;
        let loaded = bridge::open_or_import_program(&self.runtime, project, &options)?;
        Ok(LoadedProgram {
            program: Program {
                runtime: self.runtime.clone(),
                project,
                handle: Some(loaded.handle),
                _project: PhantomData,
            },
            info: loaded.info,
        })
    }

    /// Closes this project handle.
    pub fn close(mut self) -> Result<()> {
        self.close_inner()
    }

    fn handle(&self) -> Result<JavaHandle> {
        self.handle
            .ok_or_else(|| Error::closed_handle("Ghidra project"))
    }

    fn close_inner(&mut self) -> Result<()> {
        if let Some(handle) = self.handle.take() {
            bridge::close_project(&self.runtime, handle)?;
        }
        Ok(())
    }
}

impl Drop for Project {
    fn drop(&mut self) {
        let _ = self.close_inner();
    }
}

impl Program<'_> {
    /// Returns the program name stored in the Ghidra project.
    pub fn name(&self) -> Result<String> {
        bridge::program_name(&self.runtime, self.handle()?)
    }

    /// Runs default analysis when the program has not already been analyzed.
    pub fn analyze(&self) -> Result<()> {
        self.analyze_with_options(AnalysisOptions::default())?;
        Ok(())
    }

    /// Runs analysis with explicit mode and monitor options.
    pub fn analyze_with_options(&self, options: AnalysisOptions) -> Result<AnalysisReport> {
        bridge::analyze_program(&self.runtime, self.project, self.handle()?, options)
    }

    /// Saves the program back into its Ghidra project.
    pub fn save(&self) -> Result<()> {
        bridge::save_program(&self.runtime, self.project, self.handle()?)
    }

    /// Lists non-external functions in program order.
    pub fn functions(&self) -> Result<Vec<FunctionDescriptor>> {
        let json = bridge::list_functions(&self.runtime, self.handle()?)?;
        parse_json(&json, "failed to parse Ghidra function descriptors")
    }

    /// Parses a Ghidra address string in this program.
    pub fn parse_address(&self, address: &str) -> Result<Address> {
        let json = bridge::parse_address(&self.runtime, self.handle()?, address)?;
        parse_json(&json, format!("failed to parse Ghidra address {address}"))
    }

    /// Returns the function whose entry point is exactly `address`.
    pub fn function_at(&self, address: &Address) -> Result<Option<FunctionDescriptor>> {
        let Some(json) = bridge::function_at(&self.runtime, self.handle()?, address.as_str())?
        else {
            return Ok(None);
        };
        parse_json(
            &json,
            format!("failed to parse Ghidra function descriptor at {address}"),
        )
    }

    /// Returns the function whose body contains `address`.
    pub fn function_containing(&self, address: &Address) -> Result<Option<FunctionDescriptor>> {
        let Some(json) =
            bridge::function_containing(&self.runtime, self.handle()?, address.as_str())?
        else {
            return Ok(None);
        };
        parse_json(
            &json,
            format!("failed to parse Ghidra function descriptor containing {address}"),
        )
    }

    /// Lists program symbols visible through Ghidra's symbol table.
    pub fn symbols(&self) -> Result<Vec<SymbolInfo>> {
        let json = bridge::symbols(&self.runtime, self.handle()?)?;
        parse_json(&json, "failed to parse Ghidra symbols")
    }

    /// Lists symbols located exactly at `address`.
    pub fn symbols_at(&self, address: &Address) -> Result<Vec<SymbolInfo>> {
        let json = bridge::symbols_at(&self.runtime, self.handle()?, address.as_str())?;
        parse_json(
            &json,
            format!("failed to parse Ghidra symbols at {address}"),
        )
    }

    /// Finds symbols by name.
    pub fn find_symbols(&self, name: &str, case_sensitive: bool) -> Result<Vec<SymbolInfo>> {
        let json = bridge::find_symbols(&self.runtime, self.handle()?, name, case_sensitive)?;
        parse_json(
            &json,
            format!("failed to parse Ghidra symbols named {name}"),
        )
    }

    /// Lists references originating from `address`.
    pub fn references_from(&self, address: &Address) -> Result<Vec<ReferenceInfo>> {
        let json = bridge::references_from(&self.runtime, self.handle()?, address.as_str())?;
        parse_json(
            &json,
            format!("failed to parse Ghidra references from {address}"),
        )
    }

    /// Lists references targeting `address`.
    pub fn references_to(&self, address: &Address) -> Result<Vec<ReferenceInfo>> {
        let json = bridge::references_to(&self.runtime, self.handle()?, address.as_str())?;
        parse_json(
            &json,
            format!("failed to parse Ghidra references to {address}"),
        )
    }

    /// Lists memory blocks in the program.
    pub fn memory_blocks(&self) -> Result<Vec<MemoryBlockInfo>> {
        let json = bridge::memory_blocks(&self.runtime, self.handle()?)?;
        parse_json(&json, "failed to parse Ghidra memory blocks")
    }

    /// Builds a basic-block control-flow graph for one function.
    pub fn control_flow_graph(&self, function: &FunctionDescriptor) -> Result<ControlFlowGraph> {
        let json =
            bridge::control_flow_graph(&self.runtime, self.handle()?, function.entry.as_str())?;
        parse_json(
            &json,
            format!(
                "failed to parse Ghidra control-flow graph for {} at {}",
                function.name, function.entry
            ),
        )
    }

    /// Builds basic-block control-flow graphs for functions in input order.
    pub fn control_flow_graphs(
        &self,
        functions: &[FunctionDescriptor],
    ) -> Result<Vec<FunctionControlFlowGraph>> {
        if functions.is_empty() {
            return Ok(Vec::new());
        }
        let addresses = functions
            .iter()
            .map(|function| function.entry.as_str())
            .collect::<Vec<_>>();
        let json = bridge::control_flow_graphs(&self.runtime, self.handle()?, &addresses)?;
        parse_json(&json, "failed to parse Ghidra control-flow graphs")
    }

    /// Returns the linked program call graph.
    pub fn call_graph(&self) -> Result<ProgramCallGraph> {
        let json = bridge::call_graph(&self.runtime, self.handle()?)?;
        parse_json(&json, "failed to parse Ghidra call graph")
    }

    /// Returns program-level symbols, functions, and recursive type metadata.
    pub fn metadata(&self) -> Result<ProgramMetadata> {
        let json = bridge::metadata(&self.runtime, self.handle()?)?;
        parse_json(&json, "failed to parse Ghidra program metadata")
    }

    /// Collects function-local listing, reference, import, export, and source-map data.
    pub fn function_summary(&self, function: &FunctionDescriptor) -> Result<FunctionSummary> {
        let json =
            bridge::function_summary(&self.runtime, self.handle()?, function.entry.as_str())?;
        parse_json(
            &json,
            format!(
                "failed to parse Ghidra function summary for {} at {}",
                function.name, function.entry
            ),
        )
    }

    /// Collects function summaries for functions in input order.
    pub fn function_summaries(
        &self,
        functions: &[FunctionDescriptor],
    ) -> Result<Vec<FunctionSummary>> {
        if functions.is_empty() {
            return Ok(Vec::new());
        }
        let addresses = functions
            .iter()
            .map(|function| function.entry.as_str())
            .collect::<Vec<_>>();
        let json = bridge::function_summaries(&self.runtime, self.handle()?, &addresses)?;
        parse_json(&json, "failed to parse Ghidra function summaries")
    }

    /// Opens a reusable decompiler session for this program.
    pub fn open_decompiler(&self) -> Result<Decompiler<'_>> {
        let handle = bridge::open_decompiler(&self.runtime, self.handle()?)?;
        Ok(Decompiler {
            runtime: self.runtime.clone(),
            handle: Some(handle),
            _program: PhantomData,
        })
    }

    /// Decompiles a function with a timeout in seconds.
    pub fn decompile_function(
        &self,
        function: &FunctionDescriptor,
        timeout: u64,
    ) -> Result<DecompileResult> {
        let decompiler = self.open_decompiler()?;
        decompiler.decompile(function, timeout)
    }

    /// Decompiles a function with explicit decompiler options.
    pub fn decompile_function_with_options(
        &self,
        function: &FunctionDescriptor,
        options: DecompileOptions,
    ) -> Result<DecompileResult> {
        let decompiler = self.open_decompiler()?;
        decompiler.decompile_with_options(function, options)
    }

    /// Returns listing instructions in a function body.
    pub fn instructions_for_function(
        &self,
        function: &FunctionDescriptor,
    ) -> Result<Vec<InstructionInfo>> {
        let json = bridge::instructions_for_function(
            &self.runtime,
            self.handle()?,
            function.entry.as_str(),
        )?;
        parse_json(
            &json,
            format!("failed to parse Ghidra instructions for {}", function.name),
        )
    }

    /// Returns listing instructions in an address range.
    pub fn instructions_in_range(&self, range: &AddressRange) -> Result<Vec<InstructionInfo>> {
        let json = bridge::instructions_in_range(
            &self.runtime,
            self.handle()?,
            range.start.as_str(),
            range.end.as_str(),
        )?;
        parse_json(
            &json,
            format!(
                "failed to parse Ghidra instructions in range {}..{}",
                range.start, range.end
            ),
        )
    }

    /// Returns the instruction located exactly at `address`.
    pub fn instruction_at(&self, address: &Address) -> Result<Option<InstructionInfo>> {
        let Some(json) = bridge::instruction_at(&self.runtime, self.handle()?, address.as_str())?
        else {
            return Ok(None);
        };
        parse_json(
            &json,
            format!("failed to parse Ghidra instruction at {address}"),
        )
    }

    /// Returns the defined data located exactly at `address`.
    pub fn data_at(&self, address: &Address) -> Result<Option<DataInfo>> {
        let Some(json) = bridge::data_at(&self.runtime, self.handle()?, address.as_str())? else {
            return Ok(None);
        };
        parse_json(&json, format!("failed to parse Ghidra data at {address}"))
    }

    /// Returns the defined data item containing `address`.
    pub fn data_containing(&self, address: &Address) -> Result<Option<DataInfo>> {
        let Some(json) = bridge::data_containing(&self.runtime, self.handle()?, address.as_str())?
        else {
            return Ok(None);
        };
        parse_json(
            &json,
            format!("failed to parse Ghidra data containing {address}"),
        )
    }

    /// Returns Ghidra's basic block count for a function.
    pub fn basic_block_count(&self, function: &FunctionDescriptor) -> Result<u64> {
        bridge::basic_block_count(&self.runtime, self.handle()?, function.entry.as_str())
    }

    /// Returns the number of instructions in a function body.
    pub fn instruction_count(&self, function: &FunctionDescriptor) -> Result<u64> {
        bridge::instruction_count(&self.runtime, self.handle()?, function.entry.as_str())
    }

    /// Returns function entry addresses that call `function`.
    pub fn callers(&self, function: &FunctionDescriptor) -> Result<Vec<Address>> {
        let json = bridge::callers(&self.runtime, self.handle()?, function.entry.as_str())?;
        parse_json(
            &json,
            format!("failed to parse Ghidra callers for {}", function.name),
        )
    }

    /// Returns function entry addresses called by `function`.
    pub fn callees(&self, function: &FunctionDescriptor) -> Result<Vec<Address>> {
        let json = bridge::callees(&self.runtime, self.handle()?, function.entry.as_str())?;
        parse_json(
            &json,
            format!("failed to parse Ghidra callees for {}", function.name),
        )
    }

    /// Returns strings referenced from a function.
    pub fn strings(&self, function: &FunctionDescriptor) -> Result<Vec<String>> {
        let json = bridge::strings(&self.runtime, self.handle()?, function.entry.as_str())?;
        parse_json(
            &json,
            format!("failed to parse Ghidra strings for {}", function.name),
        )
    }

    /// Returns scalar constants referenced from a function.
    pub fn constants(&self, function: &FunctionDescriptor) -> Result<Vec<String>> {
        let json = bridge::constants(&self.runtime, self.handle()?, function.entry.as_str())?;
        parse_json(
            &json,
            format!("failed to parse Ghidra constants for {}", function.name),
        )
    }

    /// Returns data addresses referenced from a function.
    pub fn data_refs(&self, function: &FunctionDescriptor) -> Result<Vec<Address>> {
        let json = bridge::data_refs(&self.runtime, self.handle()?, function.entry.as_str())?;
        parse_json(
            &json,
            format!(
                "failed to parse Ghidra data references for {}",
                function.name
            ),
        )
    }

    /// Returns external symbols referenced from a function.
    pub fn imports(&self, function: &FunctionDescriptor) -> Result<Vec<ImportInfo>> {
        let json = bridge::imports(&self.runtime, self.handle()?, function.entry.as_str())?;
        parse_json(
            &json,
            format!("failed to parse Ghidra imports for {}", function.name),
        )
    }

    /// Returns exported entry-point names for a function.
    pub fn exports(&self, function: &FunctionDescriptor) -> Result<Vec<ExportInfo>> {
        let json = bridge::exports(&self.runtime, self.handle()?, function.entry.as_str())?;
        parse_json(
            &json,
            format!("failed to parse Ghidra exports for {}", function.name),
        )
    }

    /// Returns best-effort source-file identity for a function.
    pub fn source_map(&self, function: &FunctionDescriptor) -> Result<Option<SourceMapInfo>> {
        let Some(json) =
            bridge::source_map(&self.runtime, self.handle()?, function.entry.as_str())?
        else {
            return Ok(None);
        };
        parse_json(
            &json,
            format!("failed to parse Ghidra source map for {}", function.name),
        )
    }

    /// Starts a program transaction.
    pub fn start_transaction(&self, description: &str) -> Result<ProgramTransaction<'_>> {
        let handle = bridge::start_transaction(&self.runtime, self.handle()?, description)?;
        Ok(ProgramTransaction {
            runtime: self.runtime.clone(),
            handle: Some(handle),
            _program: PhantomData,
        })
    }

    /// Returns the plate comment for a function.
    pub fn function_plate_comment(&self, function: &FunctionDescriptor) -> Result<Option<String>> {
        bridge::function_plate_comment(&self.runtime, self.handle()?, function.entry.as_str())
    }

    /// Closes this program handle.
    pub fn close(mut self) -> Result<()> {
        self.close_inner()
    }

    fn handle(&self) -> Result<JavaHandle> {
        self.handle
            .ok_or_else(|| Error::closed_handle("Ghidra program"))
    }

    fn close_inner(&mut self) -> Result<()> {
        if let Some(handle) = self.handle.take() {
            bridge::close_program(&self.runtime, self.project, handle)?;
        }
        Ok(())
    }
}

impl Drop for Program<'_> {
    fn drop(&mut self) {
        let _ = self.close_inner();
    }
}

impl ProgramTransaction<'_> {
    /// Sets a function plate comment inside this transaction.
    pub fn set_function_plate_comment(
        &self,
        function: &FunctionDescriptor,
        comment: &str,
    ) -> Result<()> {
        bridge::set_function_plate_comment(
            &self.runtime,
            self.handle()?,
            function.entry.as_str(),
            comment,
        )
    }

    /// Commits this transaction.
    pub fn commit(mut self) -> Result<()> {
        self.end(true)
    }

    /// Rolls this transaction back.
    pub fn rollback(mut self) -> Result<()> {
        self.end(false)
    }

    fn handle(&self) -> Result<JavaHandle> {
        self.handle
            .ok_or_else(|| Error::closed_handle("Ghidra transaction"))
    }

    fn end(&mut self, commit: bool) -> Result<()> {
        if let Some(handle) = self.handle {
            bridge::end_transaction(&self.runtime, handle, commit)?;
            self.handle = None;
        }
        Ok(())
    }
}

impl Drop for ProgramTransaction<'_> {
    fn drop(&mut self) {
        let _ = self.end(false);
    }
}

impl Decompiler<'_> {
    /// Returns program metadata using the decompiler session type registry.
    pub fn program_metadata(&self) -> Result<ProgramMetadata> {
        let json = bridge::decompiler_metadata(&self.runtime, self.handle()?)?;
        parse_json(&json, "failed to parse Ghidra decompiler program metadata")
    }

    /// Decompiles a function with a timeout in seconds.
    pub fn decompile(
        &self,
        function: &FunctionDescriptor,
        timeout: u64,
    ) -> Result<DecompileResult> {
        let json = bridge::decompile_function(
            &self.runtime,
            self.handle()?,
            function.entry.as_str(),
            DecompileOptions::new(timeout),
        )?;
        parse_json(
            &json,
            format!(
                "failed to parse Ghidra decompile result for {} at {}",
                function.name, function.entry
            ),
        )
    }

    /// Decompiles a function with explicit decompiler options.
    pub fn decompile_with_options(
        &self,
        function: &FunctionDescriptor,
        options: DecompileOptions,
    ) -> Result<DecompileResult> {
        let json = bridge::decompile_function(
            &self.runtime,
            self.handle()?,
            function.entry.as_str(),
            options,
        )?;
        parse_json(
            &json,
            format!(
                "failed to parse Ghidra decompile result for {} at {}",
                function.name, function.entry
            ),
        )
    }

    /// Closes this decompiler session.
    pub fn close(mut self) -> Result<()> {
        self.close_inner()
    }

    fn handle(&self) -> Result<JavaHandle> {
        self.handle
            .ok_or_else(|| Error::closed_handle("Ghidra decompiler"))
    }

    fn close_inner(&mut self) -> Result<()> {
        if let Some(handle) = self.handle.take() {
            bridge::close_decompiler(&self.runtime, handle)?;
        }
        Ok(())
    }
}

impl Drop for Decompiler<'_> {
    fn drop(&mut self) {
        let _ = self.close_inner();
    }
}

fn parse_json<T, C>(json: &str, context: C) -> Result<T>
where
    T: serde::de::DeserializeOwned,
    C: JsonContext,
{
    serde_json::from_str(json).map_err(|source| context.error(source))
}

trait JsonContext {
    fn error(self, source: serde_json::Error) -> Error;
}

impl JsonContext for &'static str {
    fn error(self, source: serde_json::Error) -> Error {
        Error::json(self, source)
    }
}

impl JsonContext for String {
    fn error(self, source: serde_json::Error) -> Error {
        Error::json(self, source)
    }
}

impl<F> JsonContext for F
where
    F: FnOnce(serde_json::Error) -> Error,
{
    fn error(self, source: serde_json::Error) -> Error {
        self(source)
    }
}