//
// Copyright 2024 CloudWeGo Authors
//
// 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.
//

package obj

import (
    `encoding/binary`
    `io`
    `unsafe`
)

type MachOHeader struct {
    Magic      uint32
    CPUType    uint32
    CPUSubType uint32
    FileType   uint32
    CmdCount   uint32
    CmdSize    uint32
    Flags      uint32
    _          uint32
}

type SegmentCommand struct {
    Cmd          uint32
    Size         uint32
    Name         [16]byte
    VMAddr       uint64
    VMSize       uint64
    FileOffset   uint64
    FileSize     uint64
    MaxProtect   uint32
    InitProtect  uint32
    SectionCount uint32
    Flags        uint32
}

type SegmentSection struct {
    Name      [16]byte
    SegName   [16]byte
    Addr      uint64
    Size      uint64
    Offset    uint32
    Align     uint32
    RelOffset uint32
    RelCount  uint32
    Flags     uint32
    _         [3]uint32
}

type Registers struct {
    RAX    uint64
    RBX    uint64
    RCX    uint64
    RDX    uint64
    RDI    uint64
    RSI    uint64
    RBP    uint64
    RSP    uint64
    R8     uint64
    R9     uint64
    R10    uint64
    R11    uint64
    R12    uint64
    R13    uint64
    R14    uint64
    R15    uint64
    RIP    uint64
    RFLAGS uint64
    CS     uint64
    FS     uint64
    GS     uint64
}

type UnixThreadCommand struct {
    Cmd    uint32
    Size   uint32
    Flavor uint32
    Count  uint32
    Regs   Registers
}

const (
    _MH_MAGIC_64 = 0xfeedfacf
    _MH_EXECUTE  = 0x02
    _MH_NOUNDEFS = 0x01
)

const (
    _CPU_TYPE_I386  = 0x00000007
    _CPU_ARCH_ABI64 = 0x01000000
)

const (
    _CPU_SUBTYPE_LIB64    = 0x80000000
    _CPU_SUBTYPE_I386_ALL = 0x00000003
)

const (
    _LC_SEGMENT_64 = 0x19
    _LC_UNIXTHREAD = 0x05
)

const (
    _VM_PROT_READ    = 0x01
    _VM_PROT_WRITE   = 0x02
    _VM_PROT_EXECUTE = 0x04
)

const (
    _S_ATTR_SOME_INSTRUCTIONS = 0x00000400
    _S_ATTR_PURE_INSTRUCTIONS = 0x80000000
)

const (
    _x86_THREAD_STATE64          = 0x04
    _x86_EXCEPTION_STATE64_COUNT = 42
)

const (
    _MACHO_SIZE      = uint32(unsafe.Sizeof(MachOHeader{}))
    _SEGMENT_SIZE    = uint32(unsafe.Sizeof(SegmentCommand{}))
    _SECTION_SIZE    = uint32(unsafe.Sizeof(SegmentSection{}))
    _UNIXTHREAD_SIZE = uint32(unsafe.Sizeof(UnixThreadCommand{}))
)

const (
    _IMAGE_SIZE = 4096
    _IMAGE_BASE = 0x04000000
)

const (
    _HDR_SIZE  = _MACHO_SIZE + _SEGMENT_SIZE * 2 + _SECTION_SIZE + _UNIXTHREAD_SIZE
    _ZERO_SIZE = (_IMAGE_SIZE - _HDR_SIZE %_IMAGE_SIZE) % _IMAGE_SIZE
)

var (
    zeroBytes = [_ZERO_SIZE]byte{}
)

func assembleMachO(w io.Writer, code []byte, base uint64, entry uint64) error {
    var p0name [16]byte
    var txname [16]byte
    var tsname [16]byte

    /* segment names */
    copy(tsname[:], "__text")
    copy(txname[:], "__TEXT")
    copy(p0name[:], "__PAGEZERO")

    /* calculate size of code */
    clen := uint64(len(code))
    hlen := uint64(_HDR_SIZE + _ZERO_SIZE)

    /* Mach-O does not allow image base at zero */
    if base == 0 {
        base = _IMAGE_BASE
        entry += _IMAGE_BASE
    }

    /* Page-0 Segment */
    p0 := SegmentCommand {
        Cmd    : _LC_SEGMENT_64,
        Size   : _SEGMENT_SIZE,
        Name   : p0name,
        VMSize : base,
    }

    /* TEXT Segment */
    text := SegmentCommand {
        Cmd          : _LC_SEGMENT_64,
        Size         : _SEGMENT_SIZE + _SECTION_SIZE,
        Name         : txname,
        VMAddr       : base,
        VMSize       : hlen + clen,
        FileSize     : hlen + clen,
        MaxProtect   : _VM_PROT_READ | _VM_PROT_WRITE | _VM_PROT_EXECUTE,
        InitProtect  : _VM_PROT_READ | _VM_PROT_EXECUTE,
        SectionCount : 1,
    }

    /* __TEXT.__text section */
    tsec := SegmentSection {
        Name    : tsname,
        SegName : txname,
        Addr    : base + hlen,
        Size    : clen,
        Offset  : uint32(hlen),
        Flags   : _S_ATTR_SOME_INSTRUCTIONS | _S_ATTR_PURE_INSTRUCTIONS,
    }

    /* UNIX Thread Metadata */
    unix := UnixThreadCommand {
        Cmd    : _LC_UNIXTHREAD,
        Size   : _UNIXTHREAD_SIZE,
        Flavor : _x86_THREAD_STATE64,
        Count  : _x86_EXCEPTION_STATE64_COUNT,
        Regs   : Registers { RIP: hlen + entry },
    }

    /* Mach-O Header */
    macho := MachOHeader {
        Magic      : _MH_MAGIC_64,
        CPUType    : _CPU_ARCH_ABI64 | _CPU_TYPE_I386,
        CPUSubType : _CPU_SUBTYPE_LIB64 | _CPU_SUBTYPE_I386_ALL,
        FileType   : _MH_EXECUTE,
        CmdCount   : 3,
        CmdSize    : _SEGMENT_SIZE * 2 + _SECTION_SIZE + _UNIXTHREAD_SIZE,
        Flags      : _MH_NOUNDEFS,
    }

    /* write the headers */
    if err := binary.Write(w, binary.LittleEndian, &macho)     ; err != nil { return err }
    if err := binary.Write(w, binary.LittleEndian, &p0)        ; err != nil { return err }
    if err := binary.Write(w, binary.LittleEndian, &text)      ; err != nil { return err }
    if err := binary.Write(w, binary.LittleEndian, &tsec)      ; err != nil { return err }
    if err := binary.Write(w, binary.LittleEndian, &unix)      ; err != nil { return err }
    if err := binary.Write(w, binary.LittleEndian, &zeroBytes) ; err != nil { return err }

    /* write the code */
    if n, err := w.Write(code); err != nil {
        return err
    } else if n != len(code) {
        return io.ErrShortWrite
    } else {
        return nil
    }
}