from binaryninja.datarender import DataRenderer
from binaryninja.function import DisassemblyTextLine
from binaryninja.enums import TypeClass, HighlightStandardColor
from binaryninja.binaryview import BinaryView
from binaryninja.architecture import InstructionTextToken
from binaryninja.types import Type
from typing import List

import sys
from PySide6.QtGui import QKeySequence
from binaryninjaui import UIActionHandler, UIAction, UIActionContext

class CodeDataRenderer(DataRenderer):
	def __init__(self):
		DataRenderer.__init__(self)
	def perform_is_valid_for_data(self, _, bv: BinaryView, addr: int, type: Type, context: List[str]) -> bool:
		sym = bv.get_symbol_at(addr)
		return sym is not None and sym.name.startswith("CODE_") and type.type_class == TypeClass.ArrayTypeClass
	def perform_get_lines_for_data(self, _, bv: BinaryView, addr: int, type: Type, prefix: List[InstructionTextToken], width: int, context: List[str]) -> List[DisassemblyTextLine]:
		end = addr + len(type)
		result: List[DisassemblyTextLine] = []
		for tokens, size in bv.disassembly_tokens(addr, bv.arch):
			if addr + size > end:
				break
			result.append(DisassemblyTextLine([*tokens], addr, color=HighlightStandardColor.RedHighlightColor))
			addr += size
		return result
	def __del__(self):
		pass

def make_code(bv: BinaryView, start: int, end: int) -> None:
	if bv.get_basic_blocks_at(start):
		return
	if end - start <= 1:
		# find the next basic block, data variable, or segment/section end
		data_var = bv.get_next_data_var_after(start)
		if data_var is not None:
			end = data_var.address
		else:
			end = bv.end
		end = min(bv.get_next_basic_block_start_after(start), end)
		seg = bv.get_segment_at(start)
		if seg is not None:
			end = min(seg.end, end)
		section_ends = [s.end for s in bv.get_sections_at(start)]
		end = min(*section_ends, end)
	bv.define_user_data_var(start, Type.array(Type.int(1, False), end-start), f"CODE_{start:08x}")

def make_code_helper(ctx: UIActionContext):
	make_code(ctx.binaryView, ctx.address, ctx.address + ctx.length)

CodeDataRenderer().register_type_specific()
UIAction.registerAction("Make Code", QKeySequence("C"))
UIActionHandler.globalActions().bindAction("Make Code", UIAction(make_code_helper))