// Copyright 2021-2026 Vector 35 Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::architecture::{BranchType, CoreArchitecture};
use crate::function::Function;
use crate::rc::*;
use binaryninjacore_sys::*;
use std::fmt;
use std::fmt::Debug;
use std::hash::{Hash, Hasher};

enum EdgeDirection {
    Incoming,
    Outgoing,
}

pub struct Edge<'a, C: 'a + BlockContext> {
    pub branch: BranchType,
    pub back_edge: bool,
    pub source: Guard<'a, BasicBlock<C>>,
    pub target: Guard<'a, BasicBlock<C>>,
}

impl<'a, C: 'a + Debug + BlockContext> Debug for Edge<'a, C> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "{:?} ({}) {:?} -> {:?}",
            self.branch, self.back_edge, &*self.source, &*self.target
        )
    }
}

pub struct EdgeContext<'a, C: 'a + BlockContext> {
    dir: EdgeDirection,
    orig_block: &'a BasicBlock<C>,
}

impl<'a, C: 'a + BlockContext> CoreArrayProvider for Edge<'a, C> {
    type Raw = BNBasicBlockEdge;
    type Context = EdgeContext<'a, C>;
    type Wrapped<'b>
        = Edge<'b, C>
    where
        'a: 'b;
}

unsafe impl<'a, C: 'a + BlockContext> CoreArrayProviderInner for Edge<'a, C> {
    unsafe fn free(raw: *mut Self::Raw, count: usize, _context: &Self::Context) {
        BNFreeBasicBlockEdgeList(raw, count);
    }

    unsafe fn wrap_raw<'b>(raw: &'b Self::Raw, context: &'b Self::Context) -> Self::Wrapped<'b> {
        let edge_target = Guard::new(
            BasicBlock::from_raw(raw.target, context.orig_block.context.clone()),
            raw,
        );
        let orig_block = Guard::new(
            BasicBlock::from_raw(
                context.orig_block.handle,
                context.orig_block.context.clone(),
            ),
            raw,
        );

        let (source, target) = match context.dir {
            EdgeDirection::Incoming => (edge_target, orig_block),
            EdgeDirection::Outgoing => (orig_block, edge_target),
        };

        Edge {
            branch: raw.type_,
            back_edge: raw.backEdge,
            source,
            target,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct PendingBasicBlockEdge {
    pub branch_type: BranchType,
    pub target: u64,
    pub arch: CoreArchitecture,
    pub fallthrough: bool,
}

impl PendingBasicBlockEdge {
    pub fn new(
        branch_type: BranchType,
        target: u64,
        arch: CoreArchitecture,
        fallthrough: bool,
    ) -> Self {
        Self {
            branch_type,
            target,
            arch,
            fallthrough,
        }
    }
}

impl From<BNPendingBasicBlockEdge> for PendingBasicBlockEdge {
    fn from(edge: BNPendingBasicBlockEdge) -> Self {
        Self {
            branch_type: edge.type_,
            target: edge.target,
            arch: unsafe { CoreArchitecture::from_raw(edge.arch) },
            fallthrough: edge.fallThrough,
        }
    }
}

impl CoreArrayProvider for PendingBasicBlockEdge {
    type Raw = BNPendingBasicBlockEdge;
    type Context = ();
    type Wrapped<'a>
        = PendingBasicBlockEdge
    where
        Self: 'a;
}

unsafe impl CoreArrayProviderInner for PendingBasicBlockEdge {
    unsafe fn free(raw: *mut Self::Raw, _count: usize, _context: &Self::Context) {
        BNFreePendingBasicBlockEdgeList(raw);
    }

    unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped<'a> {
        PendingBasicBlockEdge::from(*raw)
    }
}

pub trait BlockContext: Clone + Sync + Send + Sized {
    type Instruction;
    type InstructionIndex: Debug + From<u64>;
    type Iter: Iterator<Item = Self::Instruction>;

    fn start(&self, block: &BasicBlock<Self>) -> Self::Instruction;
    fn iter(&self, block: &BasicBlock<Self>) -> Self::Iter;
}

#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy, PartialOrd, Ord)]
pub enum BasicBlockType {
    Native,
    LowLevelIL,
    MediumLevelIL,
    HighLevelIL,
}

pub struct BasicBlock<C: BlockContext> {
    pub(crate) handle: *mut BNBasicBlock,
    context: C,
}

impl<C: BlockContext> BasicBlock<C> {
    pub unsafe fn from_raw(handle: *mut BNBasicBlock, context: C) -> Self {
        Self { handle, context }
    }

    pub(crate) unsafe fn ref_from_raw(handle: *mut BNBasicBlock, context: C) -> Ref<Self> {
        Ref::new(Self::from_raw(handle, context))
    }

    // TODO native bb vs il bbs
    pub fn function(&self) -> Ref<Function> {
        unsafe {
            let func = BNGetBasicBlockFunction(self.handle);
            Function::ref_from_raw(func)
        }
    }

    pub fn arch(&self) -> CoreArchitecture {
        unsafe {
            let arch = BNGetBasicBlockArchitecture(self.handle);
            CoreArchitecture::from_raw(arch)
        }
    }

    pub fn block_type(&self) -> BasicBlockType {
        if unsafe { !BNIsILBasicBlock(self.handle) } {
            BasicBlockType::Native
        } else if unsafe { BNIsLowLevelILBasicBlock(self.handle) } {
            BasicBlockType::LowLevelIL
        } else if unsafe { BNIsMediumLevelILBasicBlock(self.handle) } {
            BasicBlockType::MediumLevelIL
        } else {
            // We checked all other IL levels, so this is safe.
            BasicBlockType::HighLevelIL
        }
    }

    pub fn iter(&self) -> C::Iter {
        self.context.iter(self)
    }

    pub fn start_index(&self) -> C::InstructionIndex {
        C::InstructionIndex::from(unsafe { BNGetBasicBlockStart(self.handle) })
    }

    pub fn end_index(&self) -> C::InstructionIndex {
        C::InstructionIndex::from(unsafe { BNGetBasicBlockEnd(self.handle) })
    }

    pub fn start(&self) -> u64 {
        unsafe { BNGetBasicBlockStart(self.handle) }
    }

    pub fn end(&self) -> u64 {
        unsafe { BNGetBasicBlockEnd(self.handle) }
    }

    pub fn set_end(&self, end: u64) {
        unsafe {
            BNSetBasicBlockEnd(self.handle, end);
        }
    }

    pub fn add_instruction_data(&self, data: &[u8]) {
        unsafe {
            BNBasicBlockAddInstructionData(self.handle, data.as_ptr() as *const _, data.len());
        }
    }

    pub fn instruction_data(&self, addr: u64) -> &[u8] {
        unsafe {
            let mut size: usize = 0;
            let data = BNBasicBlockGetInstructionData(self.handle, addr, &mut size);
            if data.is_null() {
                return &[];
            }

            std::slice::from_raw_parts(data, size)
        }
    }

    pub fn set_has_invalid_instructions(&self, value: bool) {
        unsafe {
            BNBasicBlockSetHasInvalidInstructions(self.handle, value);
        }
    }

    pub fn has_invalid_instructions(&self) -> bool {
        unsafe { BNBasicBlockHasInvalidInstructions(self.handle) }
    }

    pub fn raw_length(&self) -> u64 {
        unsafe { BNGetBasicBlockLength(self.handle) }
    }

    pub fn incoming_edges(&self) -> Array<Edge<'_, C>> {
        unsafe {
            let mut count = 0;
            let edges = BNGetBasicBlockIncomingEdges(self.handle, &mut count);
            Array::new(
                edges,
                count,
                EdgeContext {
                    dir: EdgeDirection::Incoming,
                    orig_block: self,
                },
            )
        }
    }

    pub fn outgoing_edges(&self) -> Array<Edge<'_, C>> {
        unsafe {
            let mut count = 0;
            let edges = BNGetBasicBlockOutgoingEdges(self.handle, &mut count);
            Array::new(
                edges,
                count,
                EdgeContext {
                    dir: EdgeDirection::Outgoing,
                    orig_block: self,
                },
            )
        }
    }

    /// Pending outgoing edges for the basic block. These are edges that have not yet been resolved.
    pub fn pending_outgoing_edges(&self) -> Array<PendingBasicBlockEdge> {
        unsafe {
            let mut count = 0;
            let edges_ptr = BNGetBasicBlockPendingOutgoingEdges(self.handle, &mut count);
            Array::new(edges_ptr, count, ())
        }
    }

    pub fn add_pending_outgoing_edge(&self, edge: &PendingBasicBlockEdge) {
        unsafe {
            BNBasicBlockAddPendingOutgoingEdge(
                self.handle,
                edge.branch_type,
                edge.target,
                edge.arch.handle,
                edge.fallthrough,
            );
        }
    }

    pub fn clear_pending_outgoing_edges(&self) {
        unsafe {
            BNClearBasicBlockPendingOutgoingEdges(self.handle);
        }
    }

    pub fn set_fallthrough_to_function(&self, value: bool) {
        unsafe {
            BNBasicBlockSetFallThroughToFunction(self.handle, value);
        }
    }

    pub fn is_fallthrough_to_function(&self) -> bool {
        unsafe { BNBasicBlockIsFallThroughToFunction(self.handle) }
    }

    // is this valid for il blocks? (it looks like up to MLIL it is)
    pub fn has_undetermined_outgoing_edges(&self) -> bool {
        unsafe { BNBasicBlockHasUndeterminedOutgoingEdges(self.handle) }
    }

    pub fn set_undetermined_outgoing_edges(&self, value: bool) {
        unsafe {
            BNBasicBlockSetUndeterminedOutgoingEdges(self.handle, value);
        }
    }

    pub fn can_exit(&self) -> bool {
        unsafe { BNBasicBlockCanExit(self.handle) }
    }

    pub fn set_can_exit(&self, value: bool) {
        unsafe {
            BNBasicBlockSetCanExit(self.handle, value);
        }
    }

    // TODO: Should we new type this? I just cant tell where the consumers of this are.
    pub fn index(&self) -> usize {
        unsafe { BNGetBasicBlockIndex(self.handle) }
    }

    pub fn immediate_dominator(&self) -> Option<Ref<Self>> {
        unsafe {
            // TODO: We don't allow the user to calculate post dominators
            let block = BNGetBasicBlockImmediateDominator(self.handle, false);
            if block.is_null() {
                return None;
            }
            Some(BasicBlock::ref_from_raw(block, self.context.clone()))
        }
    }

    pub fn dominators(&self) -> Array<BasicBlock<C>> {
        unsafe {
            let mut count = 0;
            // TODO: We don't allow the user to calculate post dominators
            let blocks = BNGetBasicBlockDominators(self.handle, &mut count, false);
            Array::new(blocks, count, self.context.clone())
        }
    }

    pub fn strict_dominators(&self) -> Array<BasicBlock<C>> {
        unsafe {
            let mut count = 0;
            // TODO: We don't allow the user to calculate post dominators
            let blocks = BNGetBasicBlockStrictDominators(self.handle, &mut count, false);
            Array::new(blocks, count, self.context.clone())
        }
    }

    pub fn dominator_tree_children(&self) -> Array<BasicBlock<C>> {
        unsafe {
            let mut count = 0;
            // TODO: We don't allow the user to calculate post dominators
            let blocks = BNGetBasicBlockDominatorTreeChildren(self.handle, &mut count, false);
            Array::new(blocks, count, self.context.clone())
        }
    }

    pub fn dominance_frontier(&self) -> Array<BasicBlock<C>> {
        unsafe {
            let mut count = 0;
            // TODO: We don't allow the user to calculate post dominators
            let blocks = BNGetBasicBlockDominanceFrontier(self.handle, &mut count, false);
            Array::new(blocks, count, self.context.clone())
        }
    }

    // TODO iterated dominance frontier
}

impl<C: BlockContext> Hash for BasicBlock<C> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.function().hash(state);
        self.block_type().hash(state);
        state.write_usize(self.index());
    }
}

impl<C: BlockContext> PartialEq for BasicBlock<C> {
    fn eq(&self, other: &Self) -> bool {
        self.function() == other.function()
            && self.index() == other.index()
            && self.block_type() == other.block_type()
    }
}

impl<C: BlockContext> Eq for BasicBlock<C> {}

impl<C: BlockContext> IntoIterator for &BasicBlock<C> {
    type Item = C::Instruction;
    type IntoIter = C::Iter;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

impl<C: BlockContext> IntoIterator for BasicBlock<C> {
    type Item = C::Instruction;
    type IntoIter = C::Iter;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

impl<C: fmt::Debug + BlockContext> fmt::Debug for BasicBlock<C> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("BasicBlock")
            .field("context", &self.context)
            .field("start_index", &self.start_index())
            .field("end_index", &self.end_index())
            .field("raw_length", &self.raw_length())
            .finish()
    }
}

impl<C: BlockContext> ToOwned for BasicBlock<C> {
    type Owned = Ref<Self>;

    fn to_owned(&self) -> Self::Owned {
        unsafe { RefCountable::inc_ref(self) }
    }
}

unsafe impl<C: BlockContext> RefCountable for BasicBlock<C> {
    unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
        Ref::new(Self {
            handle: BNNewBasicBlockReference(handle.handle),
            context: handle.context.clone(),
        })
    }

    unsafe fn dec_ref(handle: &Self) {
        BNFreeBasicBlock(handle.handle);
    }
}

impl<C: BlockContext> CoreArrayProvider for BasicBlock<C> {
    type Raw = *mut BNBasicBlock;
    type Context = C;
    type Wrapped<'a>
        = Guard<'a, BasicBlock<C>>
    where
        C: 'a;
}

unsafe impl<C: BlockContext> CoreArrayProviderInner for BasicBlock<C> {
    unsafe fn free(raw: *mut Self::Raw, count: usize, _context: &Self::Context) {
        BNFreeBasicBlockList(raw, count);
    }

    unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, context: &'a Self::Context) -> Self::Wrapped<'a> {
        Guard::new(BasicBlock::from_raw(*raw, context.clone()), context)
    }
}

unsafe impl<C: BlockContext> Send for BasicBlock<C> {}
unsafe impl<C: BlockContext> Sync for BasicBlock<C> {}