#!/usr/bin/env python3

from elftools.elf.elffile import ELFFile
from elftools.elf.relocation import RelocationSection
import struct, sys, itertools
from dataclasses import dataclass
from collections import namedtuple

def bitsize(i):
	return max(len(bin(i)[2:]), 1)

@dataclass
class X:
	id: int
	children: "any type"

def dump(x):
	data = b''.join(dump(c) for c in x.children) if type(x.children) == list else x.children
	l = len(data)
	l = l | (1 << ((bitsize(l) + 6) // 7 * 7))
	return x.id.to_bytes(length = (bitsize(x.id) + 7) // 8) + l.to_bytes(length = (bitsize(l) + 7) // 8) + data

ARGS = {}

for s in sys.argv[3:]:
	k, v = s.split("=", 1)
	ARGS[k] = v

SYMTRANSLATION = {}

HEADER = X(0x1A45DFA3, [X(0x4282, b'BDexeF')])
OUTPUT = X(0x10000000, [])

@dataclass
class ELFReloc:
	offset: int
	reference: int

@dataclass
class ELFSegment:
	data: bytearray
	vaddr: int
	relocations: list[ELFReloc]

with open(sys.argv[1], "rb") as inp, open(sys.argv[2], "wb") as outp:
	elf = ELFFile(inp)
	
	segments = []
	
	sections_base = min([sec["sh_offset"] for sec in elf.iter_sections() if sec.name])
	
	base_vaddr = min([seg["p_vaddr"] for seg in elf.iter_segments()])
	
	for seg in elf.iter_segments():
		if seg["p_type"] != "PT_LOAD":
			continue
		
		# If no section is in the segment, then ignore it
		if seg["p_offset"] < sections_base and seg["p_offset"] + seg["p_filesz"] < sections_base:
			continue
		
		inp.seek(seg["p_offset"])
		
		segments.append(ELFSegment(vaddr = seg["p_vaddr"], data = bytearray(inp.read(seg["p_filesz"]) + b'\0' * (seg["p_memsz"] - seg["p_filesz"])), relocations = []))
	
	rel_section = elf.get_section_by_name(".rel.dyn")
	
	if rel_section:
		for rel in rel_section.iter_relocations():
			assert rel["r_info"] % 256 == 8
			
			# The segment in which the relocation is found/required
			seg_in = next((s for s in reversed(segments) if s.vaddr <= rel["r_offset"]), None)
			
			offset = rel["r_offset"] - seg_in.vaddr
			
			val = struct.unpack("I", seg_in.data[offset : offset + 4])[0]
			
			# The segment pointed to by the relocation
			seg_ref = next((s for s in reversed(segments) if s.vaddr <= val), None)
			
			seg_in.data[offset : offset + 4] = struct.pack("I", val - seg_ref.vaddr)
			
			seg_in.relocations.append(ELFReloc(offset = rel["r_offset"] - seg_in.vaddr, reference = segments.index(seg_ref)))
	
	for seg in segments:
		OUTPUT.children.append(X(0x6200, [ #Segment
			X(0x6201, seg.data), #Segment data
			X(0x6202, b''.join(struct.pack("IB", rel.offset, rel.reference) for rel in seg.relocations)) #Segment relocations
		]))
	
	for pub, sym in [(False, s) for s in elf.get_section_by_name(".symtab").iter_symbols()] + [(True, s) for s in elf.get_section_by_name(".dynsym").iter_symbols()]:
		if not sym.name:
			continue
		
		seg = next((s for s in reversed(segments) if s.vaddr <= sym["st_value"]), None)
		
		if not seg:
			continue
		
		OUTPUT.children.append(X(0xC0, [
			X(0xC1, sym.name.encode()),
			X(0xC2, struct.pack("IB", sym["st_value"] - seg.vaddr, segments.index(seg) | (128 if pub else 0)))
		]))
	
	outp.write(dump(HEADER))
	outp.write(dump(OUTPUT))