#!/usr/bin/env python3
# SPDX-License-Identifier: MIT
import sys, pathlib, time
sys.path.append(str(pathlib.Path(__file__).resolve().parents[1]))

from m1n1.setup import *
from m1n1 import asm

from m1n1.hw.dart import DART, DARTRegs, R_TCR
from m1n1.fw import pmp

dart_addr = u.adt["arm-io/dart-pmp"].get_reg(0)[0]
pmp_dart = DART(iface, DARTRegs(u, dart_addr), u, iova_range=(0x4000, 0x1000000))

pmp_dart.initialize()

pmp_addr = u.adt["arm-io/pmp"].get_reg(0)[0]
pmpc = pmp.PMPClient(u, pmp_addr, pmp_dart)

# devs
# 5c - JPG
# 5d - msr
# 63 - gfx
# 65 - venc
# 66 - avd
# 68 - ane

## messages
#[cpu1] [ASCTracer@/arm-io/pmp] <ep:20 0000000000000000 ()
#[cpu0] [ASCTracer@/arm-io/pmp] >ep:20 000100000003c000 () configure
#[cpu0] [ASCTracer@/arm-io/pmp] <ep:20 000200000023b549 () initialized
#[cpu2] [ASCTracer@/arm-io/pmp] >ep:20 0020000000010003 ()
#[cpu0] [ASCTracer@/arm-io/pmp] <ep:20 0020100000010003 ()
#[cpu3] [ASCTracer@/arm-io/pmp] >ep:20 0020200000010000 ()
#[cpu3] [ASCTracer@/arm-io/pmp] <ep:20 0020300000010000 ()
#[cpu0] [ASCTracer@/arm-io/pmp] >ep:20 0020e00000630001 ()
#[cpu0] [ASCTracer@/arm-io/pmp] >ep:20 0020800000630000 ()
#[cpu3] [ASCTracer@/arm-io/pmp] <ep:20 0020900000630000 ()
#[cpu1] [ASCTracer@/arm-io/pmp] >ep:20 00100000003bc002 ()
#[cpu1] [ASCTracer@/arm-io/pmp] <ep:20 00110000003bc002 ()
#[cpu0] [ASCTracer@/arm-io/pmp] >ep:20 0020e00000630001 ()
#[cpu0] [ASCTracer@/arm-io/pmp] >ep:20 0020e00000660001 ()
#[cpu1] [ASCTracer@/arm-io/pmp] >ep:20 0020e00000680001 () - set state (async)
#[cpu1] [ASCTracer@/arm-io/pmp] >ep:20 0020e000005d0001 ()
#[cpu3] [ASCTracer@/arm-io/pmp] >ep:20 0020e00000660000 ()
#[cpu1] [ASCTracer@/arm-io/pmp] >ep:20 0020e00000660001 ()
#[cpu0] [ASCTracer@/arm-io/pmp] >ep:20 0020800000680000 () - set state with ack
#[cpu3] [ASCTracer@/arm-io/pmp] <ep:20 0020900000680000 () - ack
#[cpu2] [ASCTracer@/arm-io/pmp] >ep:20 0020e00000660000 ()
#[cpu0] [ASCTracer@/arm-io/pmp] >ep:20 0020e000005c0001 ()
#[cpu0] [ASCTracer@/arm-io/pmp] >ep:20 0020e000005c0000 ()
#[cpu1] [ASCTracer@/arm-io/pmp] >ep:20 0020f00000010000 () - ATC related?
#[cpu1] [ASCTracer@/arm-io/pmp] >ep:20 0020f10000010000 ()
#[cpu0] [ASCTracer@/arm-io/pmp] >ep:20 0020f00000020000 ()
#[cpu0] [ASCTracer@/arm-io/pmp] >ep:20 0020f10000020000 ()
#[cpu6] [ASCTracer@/arm-io/pmp] >ep:20 0020e00000e30001 ()
#[cpu6] [ASCTracer@/arm-io/pmp] >ep:20 0020e000015c0001 ()
#[cpu2] [ASCTracer@/arm-io/pmp] >ep:20 0020e00000e30000 ()
#[cpu2] [ASCTracer@/arm-io/pmp] >ep:20 0020e000015c0000 ()
#[cpu5] [ASCTracer@/arm-io/pmp] >ep:20 0020e000005c0001 ()
#[cpu5] [ASCTracer@/arm-io/pmp] >ep:20 0020e000005c0000 ()
#[cpu7] [ASCTracer@/arm-io/pmp] >ep:20 0020e000005c0001 ()
#[cpu7] [ASCTracer@/arm-io/pmp] >ep:20 0020e000005c0000 ()
#[cpu7] [ASCTracer@/arm-io/pmp] >ep:20 0020e000005c0001 ()
#[cpu7] [ASCTracer@/arm-io/pmp] >ep:20 0020e000005c0000 ()
#[cpu4] [ASCTracer@/arm-io/pmp] >ep:20 0020e00000650001 ()
#[cpu4] [ASCTracer@/arm-io/pmp] >ep:20 0020e00000650000 ()
#[cpu4] [ASCTracer@/arm-io/pmp] >ep:20 0020e000005d0000 ()

#[cpu7] [PMPTracer@/arm-io/pmp] [pmpeptracer] >0x20(None) 0020a00000540022 (TYPE=0x20) - add dev?
#[cpu4] [PMPTracer@/arm-io/pmp] [pmpeptracer] <0x20(None) 0020b00000540022 (TYPE=0x20) - add dev?

for i in range(10):
    pmpc.work()

LOOPS = 10000000
freq = u.mrs(CNTFRQ_EL0)

CREG = [
    0x210e00000,
    0x211e00000,
]

CLUSTER_PSTATE = 0x20020

# e-core pstates
#   600 972 1332 1704 2064
# p-core pstates
#   600 828 1056 1284 1500 1728 1956 2184 2388 2592 2772 2988 3096 3144 3204

code = u.malloc(0x1000)

util = asm.ARMAsm("""
bench:
    mrs x1, CNTPCT_EL0
1:
    sub x0, x0, #1
    cbnz x0, 1b

    mrs x2, CNTPCT_EL0
    sub x0, x2, x1
    ret
""", code)
iface.writemem(code, util.data)
p.dc_cvau(code, len(util.data))
p.ic_ivau(code, len(util.data))

def bench_cpu(idx):
    if idx == 0:
        elapsed = p.call(util.bench, LOOPS) / freq
    else:
        elapsed = p.smp_call_sync(idx, util.bench, LOOPS) / freq
    if elapsed == 0:
        return 0
    mhz = (LOOPS / elapsed) / 1000000
    return mhz

ECORE_PSTATE = 0x210e20020
PCORE_PSTATE = 0x211e20020

print()

e_pstate = p.read64(ECORE_PSTATE)
p_pstate = p.read64(PCORE_PSTATE)

print(f"E-Core pstate: {e_pstate:x}")
print(f"P-Core pstate: {p_pstate:x}")

PMGR = 0x23b700000

print("Init pokes")

if p.read64(0x210e49000) == 0:
    p.write64(0x210e49008, 0xfe527fffffffffff)
    p.write64(0x210e49000, 1)

mcc_init = [
    (0x20020058c, 0x3ff001a ),
    (0x200200590, 0x0       ),
    (0x200200594, 0x33      ),
    (0x200200598, 0x0       ),
    (0x20020059c, 0x0       ),
    (0x2002005a0, 0x0       ),
    (0x200200580, 0x71      ),
    (0x200200584, 0x1       ),
    (0x2002005a0, 0x3ff0039 ),
    (0x2002005a4, 0x0       ),
    (0x2002005a8, 0x32      ),
    (0x2002005ac, 0x0       ),
    (0x2002005b0, 0x0       ),
    (0x2002005b4, 0x0       ),
    (0x200200580, 0x71      ),
    (0x200200584, 0x3       ),
    (0x2002005b4, 0x3ff003a ),
    (0x2002005b8, 0x0       ),
    (0x2002005bc, 0x32      ),
    (0x2002005c0, 0x0       ),
    (0x2002005c4, 0x0       ),
    (0x2002005c8, 0x0       ),
    (0x200200580, 0x71      ),
    (0x200200584, 0x7       ),
    (0x2002005c8, 0x3ff0013 ),
    (0x2002005cc, 0x0       ),
    (0x2002005d0, 0x0       ),
    (0x2002005d4, 0x0       ),
    (0x2002005d8, 0x0       ),
    (0x2002005dc, 0x200     ),
    (0x200200580, 0x71      ),
    (0x200200584, 0xf       ),
    (0x2002005dc, 0x3ff0016 ),
    (0x2002005e0, 0x0       ),
    (0x2002005e4, 0x0       ),
    (0x2002005e8, 0x0       ),
    (0x2002005ec, 0x0       ),
    (0x2002005f0, 0x200     ),
    (0x200200580, 0x71      ),
    (0x200200584, 0x1f      ),
    (0x2002005f0, 0x3ff0012 ),
    (0x2002005f4, 0x0       ),
    (0x2002005f8, 0x33      ),
    (0x2002005fc, 0x0       ),
    (0x200200600, 0x0       ),
    (0x200200604, 0x200     ),
    (0x200200580, 0x71      ),
    (0x200200584, 0x3f      ),
    (0x200200604, 0x3ff0014 ),
    (0x200200608, 0x0       ),
    (0x20020060c, 0x33      ),
    (0x200200610, 0x0       ),
    (0x200200614, 0x0       ),
    (0x200200618, 0x200     ),
    (0x200200580, 0x71      ),
    (0x200200584, 0x7f      ),
    (0x200200618, 0x3ff001a ),
    (0x20020061c, 0x0       ),
    (0x200200620, 0x31      ),
    (0x200200624, 0x0       ),
    (0x200200628, 0x0       ),
    (0x20020062c, 0x0       ),
    (0x200200580, 0x71      ),
    (0x200200584, 0xff      ),
    (0x20020062c, 0x3ff001a ),
    (0x200200630, 0x0       ),
    (0x200200634, 0x32      ),
    (0x200200638, 0x0       ),
    (0x20020063c, 0x0       ),
    (0x200200640, 0x0       ),
    (0x200200580, 0x71      ),
    (0x200200584, 0x1ff     ),
]

print("MCC init seq")

for addr, val in mcc_init:
    p.write32(addr, val)

p.write32(0x23d2b000c, 0)

for cluster in range(2):
    p.write64(CREG[cluster] + 0x48000, 0)

pmgr_init = [
    (0x23d2b4100, 0x1 ),
    (0x23d2b4120, 0x1 ),
    (0x23d2b4140, 0x1 ),
    (0x23d2b4150, 0x1 ),
    (0x23d2b4160, 0x1 ),
    (0x23d2b4170, 0x1 ),
    (0x23d2b41e0, 0x1 ),
    (0x23d2b4180, 0x1 ),
    (0x23d2b41b0, 0x1 ),
    (0x23d2b41c0, 0x1 ),
    (0x23d2b41f0, 0x1 ),
    (0x23d2b41a0, 0x1 ),
    (0x23d2b41d0, 0x1 ),
    (0x23d2b4190, 0x1 ),
    (0x23d2b4210, 0x1 ),
    (0x23d2b4130, 0x1 ),
    (0x23d2b4100, 0x1 ),
    (0x23d2b4100, 0x1 ),
    (0x23d2b4200, 0x1 ),
    (0x23d2b4110, 0x1 ),
    (0x23d2b4220, 0x1 ),
    (0x23d2b4230, 0x1 ),
    (0x23d2b4240, 0x1 ),
    (0x23d2b4250, 0x1 ),
    (0x23d2b4260, 0x1 ),
    (0x23d2b4270, 0x1 ),
    (0x23d2b4280, 0x1 ),
    (0x23d2b4290, 0x1 ),
    (0x23d2b42a0, 0x1 ),
    (0x23d2b42b0, 0x1 ),
    (0x23d2b42d0, 0x1 ),
    (0x23d2b42e0, 0x1 ),
    (0x23d2b42f0, 0x1 ),
    (0x23d2b4300, 0x1 ),
    (0x23d2b42c0, 0x1 ),
    (0x23d2b4310, 0x1 ),
    (0x23d2b4320, 0x3 ),
    (0x23d2b4330, 0x1 ),
    (0x23d2b4340, 0x3 ),
    (0x23d2b4350, 0x1 ),
    (0x23d2b4360, 0x3 ),
    (0x23d2b4370, 0x1 ),
    (0x23d2b4380, 0x3 ),
    (0x23d2b4390, 0x1 ),
    (0x23d2b43a0, 0x3 ),
    (0x23d2b4000, 0x30001 ),
]

print("Pmgr init seq")

for addr, val in pmgr_init:
    p.write32(addr, val)

p.set32(PMGR + 0x18000, 0xf)
p.set32(PMGR + 0x18004, 0xf)
p.set32(PMGR + 0x1c000, 0x80000000)
p.set32(PMGR + 0x1c040, 0x80000000)

cl_init = [
    (0x210e48c10, 0x0               ),
    (0x210e48c18, 0x0               ),
    (0x210e48c90, 0x0               ),
    (0x210e48c98, 0x0               ),
    (0x210e480c0, 0x313180614c3f3131),
    (0x210e480c8, 0x3131313131313131),
    (0x210e480e0, 0x6060967f68606060),
    (0x210e480e8, 0x6060606060606060),
    (0x210e48480, 0x49490c111a2749ff),
    (0x210e48488, 0x4949494949494949),
    (0x210e48c00, 0x0               ),
    (0x210e48c08, 0x0               ),
    (0x210e48c80, 0x28280c0e121828ff),
    (0x210e48c88, 0x2828282828282828),
    (0x210e48cc0, 0x14140607090c14ff),
    (0x210e48cc8, 0x1414141414141414),
    (0x211e48c10, 0x0               ),
    (0x211e48c18, 0x0               ),
    (0x211e48c90, 0x0               ),
    (0x211e48c98, 0x0               ),
    (0x211e480c0, 0x7d71675b534a4040),
    (0x211e480c8, 0xd1d1d1d1b9a79082),
    (0x211e480e0, 0x7b726d6965656565),
    (0x211e480e8, 0xc1c1c1c1b7a69688),
    (0x211e48480, 0x333d4a5e78a3f5ff),
    (0x211e48488, 0x171718191d21272d),
    (0x211e484c0, 0x171b1b211f202700),
    (0x211e484c8, 0x141717191a      ),
    (0x211e484d0, 0x171b1b1e1f202100),
    (0x211e484d8, 0x1414161819      ),
    (0x211e484e0, 0xf14141a1b1b1d00 ),
    (0x211e484e8, 0xc0d0e0e         ),
    (0x211e48c00, 0x1214171b212a3aff),
    (0x211e48c08, 0xc0c0c0d0d0e0f10 ),
    (0x211e48c80, 0x0               ),
    (0x211e48c88, 0x0               ),
    (0x211e48cc0, 0xb0c0e10141923ff ),
    (0x211e48cc8, 0x70708080809090a ),
]

print("Cluster init seq")

for addr, val in cl_init:
    p.write64(addr, val)

for addr in [
    0x23d280058, 0x23d280060, 0x23d280068,
    0x23d280070, 0x23d2800a8, 0x23d2800b0,
    0x23b7002f8, 0x23b7003d8, 0x23b7003e0,
    0x23b700118, 0x23b700120, 0x23b700128,
    0x23b700138, 0x23b700140, 0x23b700170,
    ]:
    p.set32(addr, 0xf)
    p.set32(addr, 0x10000000)

cluster_clr = [
    0x210e20400,
    0x210e20480,
    0x210e20408,
    0x210e20488,
    0x210e20410,
    0x210e20490,
    0x210e20418,
    0x210e20498,
    0x210e20420,
    0x210e204a0,
    0x210e20428,
    0x210e204a8,
    0x211e20400,
    0x211e20480,
    0x211e20408,
    0x211e20488,
    0x211e20410,
    0x211e20490,
    0x211e20418,
    0x211e20498,
    0x211e20420,
    0x211e204a0,
    0x211e20428,
    0x211e204a8,
    0x211e20430,
    0x211e204b0,
    0x211e20438,
    0x211e204b8,
    0x211e20440,
    0x211e204c0,
    0x211e20448,
    0x211e204c8,
    0x211e20450,
    0x211e204d0,
    0x211e20458,
    0x211e204d8,
    0x211e20460,
    0x211e204e0,
    0x211e20468,
    0x211e204e8,
    0x211e20470,
    0x211e204f0,
    0x211e20478,
    0x211e204f8,
]

for off in cluster_clr:
    p.write64(off, 0)

cluster_clr = [
    0x210e20580,
    0x210e20588,
    0x210e205e0,
    0x210e205e8,
    0x210e20600,
    0x210e20608,
    0x210e20620,
    0x210e20628,
    0x210e20640,
    0x210e20648,
    0x211e20580,
    0x211e20588,
    0x211e205e0,
    0x211e205e8,
    0x211e20600,
    0x211e20608,
    0x211e20620,
    0x211e20628,
    0x211e20640,
    0x211e20648,
]

for off in cluster_clr:
    p.write64(off, 0)

for cluster in range(2):
    print(f"Initializing cluster {cluster} (early)")
    
    p.write64(CREG[cluster] + 0x20660, 0x1000000015)
    p.write64(CREG[cluster] + 0x48000, 0)
    p.write64(CREG[cluster] + 0x48080, 0xa000000000000000)

    p.clear64(CREG[cluster] + CLUSTER_PSTATE, 1<<22)
    #print("..")
    #if cluster == 1:
        #p.set64(CREG[cluster] + 0x20048, 1)

p.set32(PMGR + 0x48000, 1)
p.set32(PMGR + 0x48c00, 1)
p.set32(PMGR + 0x48800, 1)
p.set32(PMGR + 0x48400, 1)

CLUSTER_DVMR = 0x206b8
CLUSTER_LIMIT2 = 0x40240
CLUSTER_LIMIT3 = 0x40250
CLUSTER_LIMIT1 = 0x48400

PMGR_CPUGATING = 0x1c080
CLUSTER_CTRL = 0x440f8
CLUSTER_PSCTRL = 0x200f8

CLUSTER_VCTRL = 0x20000

CLUSTER_PSINFO1_SET = 0x70210
CLUSTER_PSINFO2_SET = 0x70218

for cluster in range(2):
    print(f"Initializing cluster {cluster}")
    #ena = (1<<63)
    #if cluster == 1:
        #ena |= (1<<32) | (1<<31)
    #if cluster == 1:
        #p.set64(CREG[cluster] + CLUSTER_DVMR, ena) # CLUSTER_DVMR
    
    p.set64(CREG[cluster] + CLUSTER_LIMIT1, 1<<63)
    p.clear64(CREG[cluster] + CLUSTER_LIMIT2, 1<<63)
    p.set64(CREG[cluster] + CLUSTER_LIMIT3, 1<<63)
    
    p.set64(CREG[cluster] + CLUSTER_PSTATE, 0)
    
    p.set32(PMGR + PMGR_CPUGATING + 8 * cluster, 1<<31)

    p.write64(CREG[cluster] + CLUSTER_CTRL, 1)
    
    p.set64(CREG[cluster] + CLUSTER_PSCTRL, 1<<40)

    pstate = p.read64(CREG[cluster] + CLUSTER_PSTATE) & 0xf

    p.write64(CREG[cluster] + CLUSTER_PSINFO1_SET, p.read64(CREG[cluster] + 0x70000 + pstate * 0x20))
    p.write64(CREG[cluster] + CLUSTER_PSINFO2_SET, p.read64(CREG[cluster] + 0x70008 + pstate * 0x20))

for lane in range(8):
    p.write32(0x20020000c + lane * 0x40000, 0x70a1100)

p.smp_start_secondaries()

print("== Initial CPU frequencies ==")

for cpu in range(8):
    print(f"CPU {cpu}: {bench_cpu(cpu):.2f} MHz")
    
if p.read64(0x211e49000) == 0:
    p.write64(0x211e49008, 0xfe527fffffffffff)
    p.write64(0x211e49000, 1)

for lane in range(8):
    p.write32(0x2002007a0 + lane * 0x40000, 0xc0000)

pmpc.boot()
pmpc.start_ep(0x20)

print("PMP initialized")

pmpc.work()

def set_pstate(cluster, pstate):
    #p.set64(CREG[cluster] + CLUSTER_VCTRL, 1 << 29)
    p.mask64(CREG[cluster] + CLUSTER_PSTATE, 0xf00f, (1<<25) | pstate | (pstate << 12))
    #p.write64(CREG[cluster] + 0x70210, p.read64(CREG[cluster] + 0x70000 + pstate * 0x20))
    #p.write64(CREG[cluster] + 0x70218, p.read64(CREG[cluster] + 0x70008 + pstate * 0x20))

    #p.clear64(CREG[cluster] + CLUSTER_VCTRL, 1 << 29)

    if cluster == 1:
        if pstate > 8:
            p0, p1 = 0x133, 0x55555340
        else:
            p0, p1 = 0x813057f, 0x1800180

        for lane in range(8):
            p.write32(0x200200dc4 + lane * 0x40000, p0)
            p.write32(0x200200dbc + lane * 0x40000, p1)

    p.write32(0x23b738004 + cluster*4, pstate)
    p.write32(0x23bc34000, 1 << cluster)
    pmpc.work()

for cluster in range(2):
    p.write64(CREG[cluster] + 0x44018, 0x810f)
    p.write64(CREG[cluster] + 0x44028, 0x8e0f)
    p.write64(CREG[cluster] + 0x44038, 0x9e0f)
    p.write64(CREG[cluster] + 0x44048, 0x800f)
    p.write64(CREG[cluster] + 0x44058, 0x860f)
    p.write64(CREG[cluster] + 0x44068, 0x9c0f)
    p.write64(CREG[cluster] + 0x44078, 0x9d0f)

for cluster in range(2):
    p.write64(CREG[cluster] + 0x44180, 0)

    p.write64(CREG[cluster] + 0x44188, 0)
    p.write64(CREG[cluster] + 0x44190, 0xffffffffffffffff)
    p.write64(CREG[cluster] + 0x44198, 0xffffffffffffffff)
    p.write64(CREG[cluster] + 0x441a0, 0)
    p.write64(CREG[cluster] + 0x441a8, 0)
    p.write64(CREG[cluster] + 0x441b0, 0xdbdcddde66651102)
    p.write64(CREG[cluster] + 0x441b8, 0)

    p.write64(CREG[cluster] + 0x44180, 0xff)

set_pstate(1, 7)

pmpc.pmpep.send(0x0020f00000010000)
pmpc.pmpep.send(0x0020f10000010000)
pmpc.pmpep.send(0x0020f00000020000)
pmpc.pmpep.send(0x0020f10000020000)

set_pstate(0, 1)
set_pstate(1, 15)

print()

p.write32(0x23b7002e8, 0x4ff)
p.write32(0x23b7002e8, 0x800004ff)
p.write32(0x23b7002e8, 0xff)

e_pstate = p.read64(ECORE_PSTATE)
p_pstate = p.read64(PCORE_PSTATE)

print(f"E-Core pstate: {e_pstate:x}")
print(f"P-Core pstate: {p_pstate:x}")

time.sleep(0.5)

print("== Final CPU frequencies ==")

for cpu in (0, 4):
    print(f"CPU {cpu}: {bench_cpu(cpu):.2f} MHz")

pmpc.work()