use proc_macro::TokenStream;
use proc_macro2::Span;
use quote::quote;
use syn::{
    parse, parse_macro_input, spanned::Spanned, Ident, ItemFn, ReturnType, Type, Visibility,
};

#[proc_macro_attribute]
pub fn entry(args: TokenStream, input: TokenStream) -> TokenStream {
    let f = parse_macro_input!(input as ItemFn);

    let valid_signature = f.sig.constness.is_none()
        && f.vis == Visibility::Inherited
        && f.sig.abi.is_none()
        && f.sig.inputs.is_empty()
        && f.sig.generics.params.is_empty()
        && f.sig.generics.where_clause.is_none()
        && f.sig.variadic.is_none()
        && match f.sig.output {
            ReturnType::Default => false,
            ReturnType::Type(_, ref ty) => matches! {**ty, Type::Never(_)},
        };

    if !valid_signature {
        return parse::Error::new(
            f.span(),
            "`#[entry]` function must have signature `[unsafe] fn() -> !`",
        )
        .to_compile_error()
        .into();
    }

    if !args.is_empty() {
        return parse::Error::new(Span::call_site(), "This attribute accepts no arguments")
            .to_compile_error()
            .into();
    }

    let tramp_ident = Ident::new(&format!("{}_trampoline", f.sig.ident), Span::call_site());
    let ident = &f.sig.ident;

    quote! {
        global_asm!("

    .macro offset_ttbr1, ttbr
    .endm

    .macro mov_q, reg, val
    .if (((\\val) >> 31) == 0 || ((\\val) >> 31) == 0x1ffffffff)
    movz \\reg, :abs_g1_s:\\val
    .else
    .if (((\\val) >> 47) == 0 || ((\\val) >> 47) == 0x1ffff)
    movz \\reg, :abs_g2_s:\\val
    .else
    movz \\reg, :abs_g3:\\val
    movk \\reg, :abs_g2_nc:\\val
    .endif
    movk \\reg, :abs_g1_nc:\\val
    .endif
    movk \\reg, :abs_g0_nc:\\val
    .endm


.section \".idmap.text\",\"ax\"
.globl el2_setup ; .align 2 ; el2_setup:
    msr SPsel, #1
    mrs x0, CurrentEL
    cmp x0, #(2 << 2)
    b.eq 1f
    mov_q x0, (((((1) << (11))) | (((1) << (20))) | (((1) << (22))) | (((1) << (28))) | (((1) << (29)))) | 0)
    msr sctlr_el1, x0
    mov w0, #(0xe11)
    isb
    ret

    1: mov_q x0, (((((1) << (4))) | (((1) << (5))) | (((1) << (11))) | (((1) << (16))) | (((1) << (18))) | (((1) << (22))) | (((1) << (23))) | (((1) << (28))) | (((1) << (29)))) | 0)
    msr sctlr_el2, x0

    mov_q x0, (((((1) << (11))) | (((1) << (20))) | (((1) << (22))) | (((1) << (28))) | (((1) << (29)))) | 0)
    msr sctlr_el1, x0


    mov x0, #(0x00000040 | 0x00000080 | 0x00000100 | 0x00000200 | 0x00000005)

    msr spsr_el2, x0
    msr elr_el2, lr
    mov w0, #(0xe12)
    eret
.type el2_setup, @function ; .size el2_setup, .-el2_setup

    .macro enable_dbg
    msr daifclr, #8
    .endm

    .macro reset_pmuserenr_el0, tmpreg
    mrs \\tmpreg, id_aa64dfr0_el1
    sbfx \\tmpreg, \\tmpreg, #8, #4
    cmp \\tmpreg, #1
    b.lt 9000f
    msr pmuserenr_el0, xzr
9000:
    .endm

    .macro tcr_set_t0sz, valreg, t0sz
    bfi \\valreg, \\t0sz, #0, #6
    .endm

    .macro tcr_compute_pa_size, tcr, pos, tmp0, tmp1
    mrs \\tmp0, ID_AA64MMFR0_EL1

    ubfx \\tmp0, \\tmp0, #0, #3
    mov \\tmp1, #0x5
    cmp \\tmp0, \\tmp1
    csel \\tmp0, \\tmp1, \\tmp0, hi
    bfi \\tcr, \\tmp0, \\pos, #3
    .endm

.globl __cpu_setup ; .align 2 ; __cpu_setup:
    tlbi vmalle1
    dsb nsh

    mov x0, #3 << 20
    msr cpacr_el1, x0
    mov x0, #1 << 12
    msr mdscr_el1, x0
    isb
    enable_dbg
    reset_pmuserenr_el0 x0
    ldr x5, =((0x00) << ((0) * 8)) | ((0x04) << ((1) * 8)) | ((0x0c) << ((2) * 8)) | ((0x44) << ((3) * 8)) | ((0xff) << ((4) * 8)) | ((0xbb) << ((5) * 8))
    msr mair_el1, x5

    mov_q x0, ((((1) << (0))) | (((1) << (2))) | (((1) << (3))) | (((1) << (4))) | (((1) << (8))) | (((1) << (12))) | (((1) << (14))) | (((1) << (15))) | (((1) << (18))) | (((1) << (21))) | (((1) << (23))) | 0 | (((1) << (26))) | ((((1) << (11))) | (((1) << (20))) | (((1) << (22))) | (((1) << (28))) | (((1) << (29)))))

    ldr x10, =(((((64)) - ((48))) << 0) | ((((64)) - ((48))) << 16)) | ((((1)) << 8) | (((1)) << 24)) | ((((1)) << 10) | (((1)) << 26)) | ((((3)) << 12) | (((3)) << 28)) | (((0)) << 14) | (((2)) << 30) | (((1)) << 36) | (((1)) << 37) | (((1)) << 22)

    mov x9, #16
    tcr_set_t0sz x10, x9
    tcr_compute_pa_size x10, #32, x5, x6
    msr tcr_el1, x10
    ret
.type __cpu_setup, @function ; .size __cpu_setup, .-__cpu_setup

.macro phys_to_ttbr, ttbr, phys
mov \\ttbr, \\phys
.endm

.globl __enable_mmu ; .align 2 ; __enable_mmu:
    mrs x2, ID_AA64MMFR0_EL1
    ubfx x2, x2, #28, 4
    cmp x2, #0x0
    b.ne __no_granule_support
    adrp x2, idmap_pg_dir
    phys_to_ttbr x1, x1
    phys_to_ttbr x2, x2
    msr ttbr0_el1, x2
    offset_ttbr1 x1
    msr ttbr1_el1, x1
    isb
    msr sctlr_el1, x0
    isb

    ic iallu
    dsb nsh
    isb
    ret
.type __enable_mmu, @function ; .size __enable_mmu, .-__enable_mmu

__no_granule_support:
    1:
        wfe
        wfi
        b 1b
    .type __no_granule_support, @function ; .size __no_granule_support, .-__no_granule_support

__primary_switch:
    adrp x1, swapper_pg_dir
    bl __enable_mmu
    ldr x8, =__primary_switched
    adrp x0, _text
    br x8
.type __primary_switch, @function ; .size __primary_switch, .-__primary_switch
    

    .macro read_ctr, reg
    mrs \\reg, ctr_el0
    nop
    .endm


    .macro phys_to_pte, pte, phys
    mov \\pte, \\phys
    .endm

    .macro adr_l, dst, sym
    adrp \\dst, \\sym
    add \\dst, \\dst, :lo12:\\sym
    .endm

    .macro dcache_line_size, reg, tmp
    read_ctr \\tmp
    ubfm \\tmp, \\tmp, #16, #19
    mov \\reg, #4
    lsl \\reg, \\reg, \\tmp
    .endm

    .globl __inval_dcache_area ; .align 2 ; __inval_dcache_area:
__dma_inv_area:
    add x1, x1, x0
    dcache_line_size x2, x3
    sub x3, x2, #1
    tst x1, x3
    bic x1, x1, x3
    b.eq 1f
    dc civac, x1
1: tst x0, x3
    bic x0, x0, x3
    b.eq 2f
    dc civac, x0
    b 3f
2: dc ivac, x0
3: add x0, x0, x2
    cmp x0, x1
    b.lo 2b
    dsb sy
    ret
    .globl __pi___inval_dcache_area; .type __pi___inval_dcache_area, %function; .set __pi___inval_dcache_area, __inval_dcache_area; .size __pi___inval_dcache_area, . - __inval_dcache_area; .type __inval_dcache_area, @function ; .size __inval_dcache_area, .-__inval_dcache_area
    


    .macro populate_entries, tbl, rtbl, index, eindex, flags, inc, tmp1, tmp2
.Lpe\\@: phys_to_pte \\tmp1, \\rtbl
    orr \\tmp1, \\tmp1, \\flags
    str \\tmp1, [\\tbl, \\index, lsl #3]
    mov \\tmp2, \\inc
    add \\rtbl, \\rtbl, \\tmp2
    add \\index, \\index, #1
    cmp \\index, \\eindex
    b.ls .Lpe\\@
    .endm

    .macro compute_indices, vstart, vend, shift, ptrs, istart, iend, count
    lsr \\iend, \\vend, \\shift
    mov \\istart, \\ptrs
    sub \\istart, \\istart, #1
    and \\iend, \\iend, \\istart
    mov \\istart, \\ptrs
    mul \\istart, \\istart, \\count
    add \\iend, \\iend, \\istart


    lsr \\istart, \\vstart, \\shift
    mov \\count, \\ptrs
    sub \\count, \\count, #1
    and \\istart, \\istart, \\count

    sub \\count, \\iend, \\istart
    .endm

    .macro map_memory, tbl, rtbl, vstart, vend, flags, phys, pgds, istart, iend, tmp, count, sv, tmp2
    add \\rtbl, \\tbl, #(1 << 12)
    mov \\sv, \\rtbl
    mov \\count, #0
    compute_indices \\vstart, \\vend, #((12 - 3) * (4 - (4 - 4)) + 3), \\pgds, \\istart, \\iend, \\count
    populate_entries \\tbl, \\rtbl, \\istart, \\iend, #(3 << 0), #(1 << 12), \\tmp, \\tmp2
    mov \\tbl, \\sv
    mov \\sv, \\rtbl


    compute_indices \\vstart, \\vend, #((12 - 3) * (4 - (1)) + 3), #(1 << (12 - 3)), \\istart, \\iend, \\count
    bic \\count, \\phys, #(1 << ((12 - 3) * (4 - (1)) + 3)) - 1
    populate_entries \\tbl, \\count, \\istart, \\iend, \\flags, #(1 << ((12 - 3) * (4 - (1)) + 3)), \\tmp, \\tmp2
    .endm
    
.pushsection \".head.text\",\"ax\"
_head:
  b stext
  .long 0

.global stext; .align 2 ; stext:
mov x21, x0
bl el2_setup
bl __create_page_tables
bl __cpu_setup
b __primary_switch
.type stext, @function ; .size stext, .-stext
.popsection

__primary_switched:
    adrp x4, __top_stack
    mov x1, x4
    mov sp, x1
    bl  clear_bss
    b   main


__create_page_tables:
    mov x28, lr

    adrp x0, swapper_pg_dir
    adrp x1, swapper_pg_end
    sub x1, x1, x0
    bl __inval_dcache_area

    adrp x0, swapper_pg_dir
    adrp x1, swapper_pg_end
    sub x1, x1, x0
1: stp xzr, xzr, [x0], #16
    stp xzr, xzr, [x0], #16
    stp xzr, xzr, [x0], #16
    stp xzr, xzr, [x0], #16
    subs x1, x1, #64
    b.ne 1b

    mov x7, (((4) << 2) | ((1 << 0) | (1 << 10) | (3 << 8)))

    adrp x0, idmap_pg_dir
    adrp x3, __idmap_text_start

    dmb sy
    dc ivac, x6
    adrp x5, __idmap_text_end
    clz x5, x5
    cmp x5, ((((64)) - ((48))) << 0)
    b.ge 1f
1:
    mov x4, #512
    mov x5, x3
    adr_l x6, __idmap_text_end

    map_memory x0, x1, x3, x6, x7, x3, x4, x10, x11, x12, x13, x14, x9

    adrp x0, swapper_pg_dir
    mov_q x5, 0xffff000000000000
    mov x4, 512
    adrp x6, _end
    adrp x3, _text
    sub x6, x6, x3
    sub x6, x6, x5

    map_memory x0, x1, x5, x6, x7, x3, x4, x10, x11, x12, x13, x14, x9

    adrp x0, idmap_pg_dir
    adrp x1, _end
    sub x1, x1, x0
    dmb sy
    bl __inval_dcache_area

    ret x28
    .type __create_page_tables, @function ; .size __create_page_tables, .-__create_page_tables
    .ltorg
    


    "
        );

        #[doc(hidden)]
        #[export_name = "main"]
        pub unsafe extern "C" fn #tramp_ident() {
            #ident()
        }

        #f
    }
    .into()
}