import struct

import idaapi
import ida_idp
import ida_loader
import ida_segment
import ida_entry
import ida_name
import ida_ua
import ida_bytes

DOL_FORMAT_NAME = "Nintendo GameCube/Wii DOL"
PPC_PROCESSOR = "ppc"

TEXT_COUNT = 7
DATA_COUNT = 11
DOL_HEADER_SIZE = 0x100


def read_u32be(li):
    b = li.read(4)
    if len(b) != 4:
        raise EOFError("Unexpected EOF")
    return struct.unpack(">I", b)[0]


def get_file_size(li):
    cur = li.tell()
    li.seek(0, 2)
    size = li.tell()
    li.seek(cur, 0)
    return size


def is_mem_addr(addr):
    # GameCube/Wii MEM1 + common Wii MEM2 range.
    return (
        0x80000000 <= addr < 0x81800000 or
        0x90000000 <= addr < 0x94000000
    )


def parse_dol(li):
    li.seek(0)

    text_offsets = [read_u32be(li) for _ in range(TEXT_COUNT)]
    data_offsets = [read_u32be(li) for _ in range(DATA_COUNT)]

    text_addrs = [read_u32be(li) for _ in range(TEXT_COUNT)]
    data_addrs = [read_u32be(li) for _ in range(DATA_COUNT)]

    text_sizes = [read_u32be(li) for _ in range(TEXT_COUNT)]
    data_sizes = [read_u32be(li) for _ in range(DATA_COUNT)]

    bss_addr = read_u32be(li)
    bss_size = read_u32be(li)
    entry = read_u32be(li)

    return {
        "text_offsets": text_offsets,
        "data_offsets": data_offsets,
        "text_addrs": text_addrs,
        "data_addrs": data_addrs,
        "text_sizes": text_sizes,
        "data_sizes": data_sizes,
        "bss_addr": bss_addr,
        "bss_size": bss_size,
        "entry": entry,
    }


def validate_dol(li):
    size = get_file_size(li)

    if size < DOL_HEADER_SIZE:
        return False

    li.seek(0)
    magic = li.read(4)

    # Reject obvious non-DOL containers/files.
    if magic == b"\x55\xAA\x38\x2D":  # U8 archive
        return False
    if magic in (b"Yaz0", b"LZ77", b"RARC", b"RIFF", b"\x89PNG"):
        return False
    if magic[:2] in (b"PK", b"BM"):
        return False

    try:
        h = parse_dol(li)
    except Exception:
        return False

    sections = []

    for i in range(TEXT_COUNT):
        off = h["text_offsets"][i]
        addr = h["text_addrs"][i]
        sz = h["text_sizes"][i]

        if sz == 0:
            continue

        if off < DOL_HEADER_SIZE or off >= size:
            return False
        if off + sz > size:
            return False
        if off % 4 != 0:
            return False
        if not is_mem_addr(addr):
            return False
        if not is_mem_addr(addr + sz - 1):
            return False

        sections.append(("text", i, off, addr, sz))

    for i in range(DATA_COUNT):
        off = h["data_offsets"][i]
        addr = h["data_addrs"][i]
        sz = h["data_sizes"][i]

        if sz == 0:
            continue

        if off < DOL_HEADER_SIZE or off >= size:
            return False
        if off + sz > size:
            return False
        if off % 4 != 0:
            return False
        if not is_mem_addr(addr):
            return False
        if not is_mem_addr(addr + sz - 1):
            return False

        sections.append(("data", i, off, addr, sz))

    if len(sections) < 2:
        return False

    # Entry must point into a text section.
    entry_ok = False
    for kind, _idx, _off, addr, sz in sections:
        if kind == "text" and addr <= h["entry"] < addr + sz:
            entry_ok = True
            break

    if not entry_ok:
        return False

    # File ranges should not overlap.
    ranges = sorted((off, off + sz) for _kind, _idx, off, _addr, sz in sections)
    for a, b in zip(ranges, ranges[1:]):
        if a[1] > b[0]:
            return False

    return True


def accept_file(li, filename):
    if validate_dol(li):
        return {
            "format": DOL_FORMAT_NAME,
            "processor": PPC_PROCESSOR,
        }

    return 0


def add_segment(start, end, name, seg_class):
    ida_segment.add_segm(0, start, end, name, seg_class)


def load_file(li, neflags, fmt):
    print("[DOL] Loading Nintendo GameCube/Wii DOL")

    ida_idp.set_processor_type(PPC_PROCESSOR, ida_idp.SETPROC_LOADER)

    h = parse_dol(li)

    for i in range(TEXT_COUNT):
        off = h["text_offsets"][i]
        addr = h["text_addrs"][i]
        sz = h["text_sizes"][i]

        if sz == 0:
            continue

        end = addr + sz
        name = ".text%d" % i

        print("[DOL] text%d file=0x%X addr=0x%X size=0x%X" % (i, off, addr, sz))

        add_segment(addr, end, name, "CODE")
        li.file2base(off, addr, end, True)

        seg = ida_segment.getseg(addr)
        if seg:
            seg.perm = ida_segment.SEGPERM_READ | ida_segment.SEGPERM_EXEC
            ida_segment.update_segm(seg)

        ida_ua.create_insn(addr)

    for i in range(DATA_COUNT):
        off = h["data_offsets"][i]
        addr = h["data_addrs"][i]
        sz = h["data_sizes"][i]

        if sz == 0:
            continue

        end = addr + sz
        name = ".data%d" % i

        print("[DOL] data%d file=0x%X addr=0x%X size=0x%X" % (i, off, addr, sz))

        add_segment(addr, end, name, "DATA")
        li.file2base(off, addr, end, True)

        seg = ida_segment.getseg(addr)
        if seg:
            seg.perm = ida_segment.SEGPERM_READ | ida_segment.SEGPERM_WRITE
            ida_segment.update_segm(seg)

    if h["bss_size"] != 0 and is_mem_addr(h["bss_addr"]):
        bss_start = h["bss_addr"]
        bss_end = h["bss_addr"] + h["bss_size"]

        print("[DOL] bss addr=0x%X size=0x%X" % (bss_start, h["bss_size"]))

        add_segment(bss_start, bss_end, ".bss", "BSS")

        seg = ida_segment.getseg(bss_start)
        if seg:
            seg.perm = ida_segment.SEGPERM_READ | ida_segment.SEGPERM_WRITE
            ida_segment.update_segm(seg)

    ida_entry.add_entry(h["entry"], h["entry"], "entry", True)
    ida_name.set_name(h["entry"], "entry", ida_name.SN_CHECK)

    print("[DOL] Entry point: 0x%08X" % h["entry"])

    return 1