package com.seaofnodes.simple;

import com.seaofnodes.simple.codegen.RegAllocTestSupport.CheckedCodeGen;
import com.seaofnodes.simple.util.Utils;
import com.seaofnodes.simple.codegen.CodeGen;
import com.seaofnodes.simple.node.cpus.arm.arm;
import com.seaofnodes.simple.node.cpus.riscv.riscv;
import java.io.IOException;
import org.junit.Ignore;
import org.junit.Test;
import static org.junit.Assert.*;

public class Chapter22Test {

    @Test public void testNarrowCReturnsEmulated() {
        String[] types = {"i8","u8","i16","u16","i32","u32","int"};
        long[] values = {0xa5a5a5a500000000L,0xa5a5a5a50000007fL,0xa5a5a5a500000080L,
            0xa5a5a5a500007fffL,0xa5a5a5a500008000L,0xa5a5a5a57fffffffL,
            0xa5a5a5a580000000L,0xa5a5a5a5ffffffffL};
        for( String target : new String[]{"arm","riscv"} )
            for( int i=0; i<types.length; i++ ) {
                String src = "{"+types[i]+"} c_narrow=\"C\"; return c_narrow();";
                CodeGen code = new CheckedCodeGen(src).driver(CodeGen.Phase.Encoding,target,"SystemV");
                var enc = code._encoding;
                // Link the external call to a tiny independent native stub.
                var call = code._stop.walk(n -> n instanceof com.seaofnodes.simple.node.CallNode ? n : null);
                int start = enc.opStart(call);
                ((com.seaofnodes.simple.codegen.RIPRelSize)call).patch(enc,start,enc.opLen(call),0x1000-start);
                for( long value : values ) {
                    long expected = switch(i) {
                    case 0 -> (byte)value; case 1 -> value&255;
                    case 2 -> (short)value; case 3 -> value&65535;
                    case 4 -> (int)value; case 5 -> value&0xffffffffL;
                    default -> value;
                    };
                    byte[] image = new byte[1<<16];
                    byte[] bits = enc.bits();
                    System.arraycopy(bits,0,image,0,bits.length);
                    if( target.equals("arm") ) {
                        EvalArm64 cpu = new EvalArm64(image,1<<15);
                        // MOVZ X0,0x1100; LDR X0,[X0]; RET X30.
                        cpu.st4(0x1000,0xd2822000); cpu.st4(0x1004,0xf9400000);
                        cpu.st4(0x1008,0xd65f03c0);
                        cpu.st8(0x1100,value);
                        assertEquals(0,cpu.step(1000));
                        assertEquals(0,cpu._pc);
                        assertEquals(types[i],expected,cpu.regs[0]);
                    } else {
                        EvalRisc5 cpu = new EvalRisc5(image,1<<15);
                        // AUIPC a0,0; LD a0,12(a0); RET.
                        cpu.st4(0x1000,0x00000517); cpu.st4(0x1004,0x00c53503);
                        cpu.st4(0x1008,0x00008067);
                        // RV64 widens even u32 returns by sign-extending bit 31.
                        cpu.st8(0x100c,i==6 ? value : (int)expected);
                        assertEquals(0,cpu.step(1000));
                        assertEquals(0,cpu._pc);
                        assertEquals(types[i],expected,cpu.regs[riscv.A0]);
                    }
                }
            }
    }

    @Test public void testNarrowCReturns() throws IOException {
        String src = """
            {i8} c_i8="C"; {u8} c_u8="C";
            {i16} c_i16="C"; {u16} c_u16="C";
            {i32} c_i32="C"; {u32} c_u32="C";
            {int} c_i64="C";
            val test_i8={->return c_i8();};
            val test_u8={->return c_u8();};
            val test_i16={->return c_i16();};
            val test_u16={->return c_u16();};
            val test_i32={->return c_i32();};
            val test_u32={->return c_u32();};
            val test_i64={->return c_i64();};
            val test_cmp={->return c_i32()==-1;};
            """;
        for( String conv : new String[]{"SystemV","win64"} )
            TestC.run(src,"c_returns"+conv,null,conv,
                      conv.equals("SystemV") ? "sysv_abi" : "ms_abi",
                      TestC.C_DRIVERS_DIR+"c_returns.c","",-1);
    }


    @Test public void testSubZeroTypeError() {
        try {
            new CodeGen("return null-0;").parse().opto().typeCheck();
            fail("Subtraction must reject null even when the other operand is zero");
        } catch( Parser.ParseException e ) {
            assertEquals("Cannot null - 0",e.getMessage());
        }
    }


    @Test
    public void testJig() throws IOException {
        String src = """
struct s0 {
    bool v1;
    i16 v2;
    int v3;
    i8 v4;
    byte v5;
};
while(new s0.v3)
    while(new s0.v5<<new s0.v4) {}
if(0) {
    if(0) {
        flt !P5ZUD4=new s0.v2;
    }
    while(0) {}
}
return new s0.v1;
""";
        testCPU(src,"x86_64_v2", "win64"  ,-1,null);
        testCPU(src,"riscv"    , "SystemV",-1,null);
        testCPU(src,"arm"      , "SystemV",-1,null);
    }

    static CodeGen testCPU( String src, String cpu, String os, int spills, String stop ) {
        CodeGen code = new CheckedCodeGen(src).driver(CodeGen.Phase.Encoding,cpu,os);
        int delta = spills>>3;
        if( delta==0 ) delta = 1;
        if( spills != -1 && !CodeGen.iterSeedOverridden() )
            assertEquals("Expect spills:",spills,code._regAlloc._spillScaled,delta);
        if( stop != null )
            assertEquals(stop, code.print());
        return code;
    }


    static int testCPUSize( String src, String cpu, String os, int spills ) {
        return testCPU(src,cpu,os,spills, "return 0;" )._encoding._bits.size();
    }

    // Should not fold away
    @Test public void testSextFail() throws IOException {
        String src = """
struct _Person { i32 age;};
_Person !p = new _Person;
p.age = (arg<<17)>>17;
return 0;
""";
        assertEquals(82, testCPUSize(src, "x86_64_v2","win64", 3 ));
        assertEquals(96, testCPUSize(src, "riscv",  "SystemV", 4 ));
        assertEquals(96, testCPUSize(src, "arm"  ,  "SystemV", 4 ));

        // do assertEquals here
        EvalRisc5 R5 = TestRisc5.build( src, "sext_str_not_fold_away", 0, 4, false);
        int trap = R5.step(100);
        assertEquals(0,trap);

        EvalArm64 A5 = TestArm64.build("sext_str_not_fold_away", src, 0, 4, false);
        trap = A5.step(100);
        assertEquals(0,trap);

    }

    // Should not fold away
    @Test public void testSextFail2() throws IOException {
        String src = """
                struct _Person { i32 age;};
                _Person !p = new _Person;
                p.age = (arg<<48)>>48;
                return 0;
        """;

        EvalRisc5 R5 = TestRisc5.build( src, "sext_str_not_fold_away_2", 0, 4, false);
        int trap = R5.step(100);
        assertEquals(0,trap);

        EvalArm64 A5 = TestArm64.build("sext_str_not_fold_away_2", src, 0, 3, false);
        int trap_arm = A5.step(100);
        assertEquals(0,trap_arm);

        assertEquals(82, testCPUSize(src, "x86_64_v2","win64",3 ));
        assertEquals(96, testCPUSize(src, "riscv",  "SystemV",4 ));
        assertEquals(96, testCPUSize(src, "arm",    "SystemV",4 ));

    }

    // Should fold away
    @Test public void testSextSuccess() throws IOException {
        String src = """
// Should fold away sign extend
struct _Person { i8 age;};
_Person !p = new _Person;
p.age = (arg<<48)>>48;
return 0;
""";

        EvalRisc5 R5 = TestRisc5.build( src, "sext_str_fold_away", 0, 5, false);
        int trap = R5.step(100);
        assertEquals(0,trap);

        EvalArm64 A5 = TestArm64.build("sext_str_fold_away", src, 0, 5, false);
        int trap_arm = A5.step(100);
        assertEquals(0,trap_arm);

        assertEquals(72, testCPUSize(src, "x86_64_v2","win64",3 ));
        assertEquals(92, testCPUSize(src, "riscv",  "SystemV",5 ));
        assertEquals(92, testCPUSize(src, "arm",    "SystemV",5));
        // do assertEquals here
    }


    // Int now is changed to 4 bytes.
    @Test public void testPerson() throws IOException {
        String src =
"""
struct Person { // Exports a field typed singleton+final *class:Person{age:i32}
    i32 age;
};

val fcn = { Person?[] ps, int x -> // exports a field with a constant fcn ptr; exports the fcn also
    if( ps[x] )
        ps[x].age++;
};
""";
        String person = "6\n";
        TestC.run(src, "person2", null, TestC.CALL_CONVENTION, "", TestC.C_DRIVERS_DIR+"person.c", person, 2);

        // Memory layout starting at PS:
        int ps = 1<<16;         // Person array pointer starts at heap start
        // Person[3] = { len,pad,P0,P1,P2 }; // sizeof = 4*8
        // P0 = { age } // sizeof=8
        int p0 = ps+4*8+0*8;
        // P1 = { age } // sizeof=8
        int p1 = ps+4*8+1*8;
        // P2 = { age } // sizeof=8
        int p2 = ps+4*8+2*8;
        EvalRisc5 R5 = TestRisc5.build( src, "fcn", ps, 2, false);
        R5.regs[riscv.A1] = 1;  // Index 1
        R5.st8(ps,3);           // Length
        R5.st8(ps+1*8,p0);
        R5.st8(ps+2*8,p1);
        R5.st8(ps+3*8,p2);
        R5.st8(p0, 5); // age= 5
        R5.st8(p1,17); // age=17
        R5.st8(p2,60); // age=60

        int trap = R5.step(100);
        assertEquals(0,trap);
        assertEquals( 5+0,R5.ld8(p0));
        assertEquals(17+1,R5.ld8(p1));
        assertEquals(60+0,R5.ld8(p2));

        EvalArm64 A5 = TestArm64.build("fcn", src, ps, 2, false);
        A5.regs[arm.X1] = 1;  // Index 1
        A5.st8(ps, 3);
        A5.st8(ps+1*8,p0);
        A5.st8(ps+2*8,p1);
        A5.st8(ps+3*8,p2);
        A5.st8(p0, 5); // age= 5
        A5.st8(p1,17); // age=17
        A5.st8(p2,60); // age=60

        int trap_arm = A5.step(100);
        assertEquals(0,trap_arm);
        assertEquals( 5+0, A5.ld8(p0));
        assertEquals(17+1, A5.ld8(p1));
        assertEquals(60+0, A5.ld8(p2));
    }

    @Test
    public void testCoRecur() {
        String src = """
val x = 5; // aa.az; // Error to self-define forward ref
struct _A { Test._B? b; Test._C? c; i64 ax; val az = x*2; new _A = { i64 x -> ax=x; }; };
struct _B { Test._A? a; Test._C? c; f32 bx; val bz = x*3; new _B = { f32 x, Test._A? aa -> bx=x; a=aa; }; };
struct _C { Test._A? a; Test._B? b; f64 cx; val cz = x*x; new _C = { f64 x, Test._A? aa, Test._B? bb -> cx=x; a=aa; b=bb; }; };
_A !aa = new _A(17);
_B !bb = new _B(3.14f, aa);
_C !cc = new _C(2.73, aa, bb);
aa.b = bb;
aa.c = cc;
bb.c = cc;
return cc.cz;
""";
        CodeGen code = new CodeGen(src).parse().opto().typeCheck();
        assertEquals("return 25;", code.print());
        assertEquals("25", Eval2.eval(code, 0));
    }

    @Test
    public void testHelloWorld() throws IOException {
        String src = """
// fd  buf len -> len
{  i32 i64 u32 -> u32 } write = "C";
// Write byte array to stdout
val p = { u8[~] str ->
    i64 ptr = str;  // cast array base to i64
    return write(1,ptr,str#);
};
p("Hello, World!");
return 0;
""";
        TestC.runSF(src,"helloWorld", "Hello, World!",2);

        // Evaluate on RISC5 emulator
        EvalRisc5 R5 = TestRisc5.build( src, "helloWorld", 0, 5, false);
        int trap = R5.step(100);
        assertEquals(0,trap);
        assertEquals(0,R5.regs[riscv.A0]);
        assertEquals("Hello, World!",R5._stdout.toString());

        // Evaluate on ARM emulator
        EvalArm64 arm = TestArm64.build("helloWorld", src,0, 5, false);
        trap = arm.step(100);
        assertEquals(0,trap);
        assertEquals(0,arm.regs[0]);
        assertEquals("Hello, World!",arm._stdout.toString());
    }


    @Test
    public void testStrLoad() {
        CodeGen code = new CodeGen("return \"abc\"[1];");
        code.parse().opto();
        assertEquals("return .[];", code.print());
    }


    @Test
    public void testFinalArray() {
        String src = """
int N=4;
i32[] !is = new i32[N];
for( int i=0; i<N; i++ )
    is[i] = i*i;
val _sum = { i32[~] is ->  // final array
    int sum=0;
    for( int i=0; i<is#; i++ )
        sum += is[i];
    return sum;
};
return _sum(is);
""";
        CodeGen code = new CodeGen(src).parse().opto().typeCheck();
        assertEquals("return Phi(Loop,0,(Phi_sum+.[]));", code.print());
        assertEquals("14", Eval2.eval(code, 0));
    }


    @Test @Ignore
    public void testEcho() throws IOException {
        String src =
"""
// Echo stdin to stdout.
return sys.io.p( sys.io.stdin() );
""";
        TestC.runSF(src,"echo", "",0);

        // Evaluate on RISC5 emulator
        EvalRisc5 R5 = TestRisc5.build( src, "echo", 0, 2, false);
        int trap = R5.step(100);
        assertEquals(0,trap);
        assertEquals(0,R5.regs[riscv.A0]);

        // Evaluate on ARM emulator
        EvalArm64 arm = TestArm64.build("echo", src, 0, 2, false);
        trap = arm.step(100);
        assertEquals(0,trap);
        assertEquals(0,arm.regs[0]);
    }


}