//
// Copyright (c) 2023 Google LLC
//
// 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
//
//     https://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.

using System;
using System.Linq;
using System.Runtime.InteropServices;
using System.Threading;
using System.Collections.Generic;
using System.Collections.Concurrent;
using Antmicro.Renode.Core;
using Antmicro.Renode.Core.Structure.Registers;
using Antmicro.Renode.Exceptions;
using Antmicro.Renode.Logging;
using Antmicro.Renode.Peripherals.Bus;
using Antmicro.Renode.Peripherals.CPU;
using Antmicro.Renode.Peripherals.Memory;
using Antmicro.Renode.Peripherals.Timers;
using Antmicro.Renode.Peripherals.CPU.Disassembler;
using Antmicro.Renode.Peripherals.CPU.Registers;
using Antmicro.Renode.Utilities;
using Antmicro.Renode.Time;
using Antmicro.Renode.Utilities.Binding;
using ELFSharp.ELF;
using ELFSharp.UImage;
using Machine = Antmicro.Renode.Core.Machine;


namespace Antmicro.Renode.Peripherals.CPU
{

    public class KelvinCPU : BaseCPU, ICPUWithRegisters, ITimeSink, IDisposable
    {
        [StructLayout(LayoutKind.Explicit, CharSet = CharSet.Ansi, Pack = 1)]
        private struct RegInfo
        {
            [FieldOffset(0)]
            public Int32 index;
            [FieldOffset(4)]
            public Int32 width;
            [FieldOffset(8)]
            public bool isGeneral;
            [FieldOffset(9)]
            public bool isReadOnly;
        }

        public KelvinCPU(uint id, string cpuType, Machine machine,
            MappedMemory memory,
            Endianess endianess = Endianess.LittleEndian,
            CpuBitness bitness = CpuBitness.Bits32)
            : base(id, cpuType, machine, endianess, bitness)
        {
            error_ptr = Marshal.AllocHGlobal(4);
            value_ptr = Marshal.AllocHGlobal(8);
            reg_info_ptr = Marshal.AllocHGlobal(Marshal.SizeOf<RegInfo>());
            string_ptr = Marshal.AllocHGlobal(maxStringLen);
            dataMemory = memory;

            // Touch all the memory segments to ensure all of segments are
            // allocated with valid pointers.
            dataMemory.TouchAllSegments();
            cpuLibraryRegistered = false;
        }

        ~KelvinCPU()
        {
            Marshal.FreeHGlobal(error_ptr);
            Marshal.FreeHGlobal(value_ptr);
            Marshal.FreeHGlobal(reg_info_ptr);
            Marshal.FreeHGlobal(string_ptr);
        }

        public override void Start()
        {
            base.Start();
            if (kelvin_id < 0)
            {
                LogAndThrowRE("Failed to create kelvin sim instance");
                return;
            }
            PC = 0;
        }

        public void RegisterControlBlock(MlTopControlBlock controlBlock)
        {
            this.controlBlock = controlBlock;
            controlBlockRegistered = true;
        }

        private void HandleKelvinEBreak()
        {
            if (controlBlockRegistered)
            {
                controlBlock.ExecEBreak();
            }
            else
            {
                IsHalted = true;
            }
        }

        private void HandleKelvinMPause()
        {
            if (controlBlockRegistered)
            {
                controlBlock.ExecMPause();
            }
            else
            {
                IsHalted = true;
            }
        }

        public override void Reset()
        {
            base.Reset();
            instructionsExecutedThisRound = 0;
            totalExecutedInstructions = 0;
            // Reset simulator.
            var result = reset(kelvin_id);
            if (result < 0)
            {
                LogAndThrowRE("Failed to reset kelvin");
            }
            this.NoisyLog("Reset kelvin");
        }

        public override void Dispose()
        {
            base.Dispose();
            // Cleanup: simulator and any unmanaged resources.
            lock(nativeLock)
            {
                destruct(kelvin_id);
            }
        }

        // No-op public API to match RiscVCPU
        public uint PerformanceInMips { get; set; }

        public override string Architecture { get { return "kelvin"; } }

        public override ulong ExecutedInstructions => totalExecutedInstructions;

        public override ExecutionMode ExecutionMode
        {
            get
            {
                return executionMode;
            }
            set
            {
                lock (singleStepSynchronizer.Guard)
                {
                    if (executionMode == value)
                    {
                        return;
                    }

                    executionMode = value;

                    singleStepSynchronizer.Enabled = IsSingleStepMode;
                    UpdateHaltedState();
                }
            }
        }

        [Register]
        public override RegisterValue PC
        {
            get
            {
                Int64 value = 0;
                lock(nativeLock)
                {
                    Int32 error = read_register(kelvin_id, PC_ID, value_ptr);
                    // TODO(hcindyl): Check different error types.
                    if (error < 0)
                    {
                        this.ErrorLog("Failed to read PC");
                    }
                    value = Marshal.ReadInt64(value_ptr);

                }
                return Convert.ToUInt64(value);
            }

            set
            {
                lock(nativeLock)
                {
                    Int32 error = write_register(kelvin_id, PC_ID, value);

                    // TODO(hcindyl): Check different error types.
                    if (error < 0)
                    {
                        this.NoisyLog("Failed to write PC");
                    }
                }
            }
        }

        public override ExecutionResult ExecuteInstructions(
            ulong numberOfInstructionsToExecute,
            out ulong numberOfExecutedInstructions)
        {
            UInt64 instructionsExecutedThisRound = 0UL;
            ExecutionResult result = ExecutionResult.Ok;
            try
            {
                lock(nativeLock)
                {
                    // Invoke simulator for the number of instructions.
                    instructionsExecutedThisRound = step(kelvin_id,
                                                         numberOfInstructionsToExecute,
                                                         error_ptr);
                    // Parse the result.
                    Int32 step_result = Marshal.ReadInt32(error_ptr);
                    switch (step_result)
                    {
                        case -1:
                            result = ExecutionResult.Aborted;
                            break;
                        case 0:
                            result = ExecutionResult.Ok;
                            break;
                        case 1:
                            result = ExecutionResult.Interrupted;
                            break;
                        case 2:
                            result = ExecutionResult.WaitingForInterrupt;
                            break;
                        case 3:
                            result = ExecutionResult.StoppedAtBreakpoint;
                            break;
                        case 4:
                            result = ExecutionResult.StoppedAtWatchpoint;
                            break;
                        case 5:
                            result = ExecutionResult.ExternalMmuFault;
                            break;
                        default:
                            LogAndThrowRE("Unknown return value from step - " + step_result);
                            break;
                    }
                }
            }
            catch (Exception)
            {
                this.NoisyLog("CPU exception detected, halting.");
                InvokeHalted(new HaltArguments(HaltReason.Abort, this));
                return ExecutionResult.Aborted;
            }
            finally
            {
                numberOfExecutedInstructions = instructionsExecutedThisRound;
                totalExecutedInstructions += instructionsExecutedThisRound;
            }

            if (numberOfInstructionsToExecute > instructionsExecutedThisRound &&
                result == ExecutionResult.Ok)
            {
                this.DebugLog("CPU finish execution.");
                this.DebugLog("Total instruction count: {0}",
                              totalExecutedInstructions);
                HandleKelvinMPause();
            }

            if (numberOfInstructionsToExecute > instructionsExecutedThisRound &&
                result == ExecutionResult.Aborted)
            {
                this.DebugLog("CPU finish execution.");
                HandleKelvinEBreak();
            }

            return result;
        }

        // ICPUWithRegisters methods implementations.
        public void SetRegister(int register, RegisterValue value)
        {
            lock(nativeLock)
            {
                var status = write_register((Int32)kelvin_id, (Int32)register, (UInt64)value);
                if (status < 0)
                {
                    LogAndThrowRE("Failed to write register " + register);
                }
            }
        }

        public RegisterValue GetRegister(int register)
        {
            Int64 value = 0;
            lock(nativeLock)
            {
                var status = read_register(kelvin_id, (UInt32)register, value_ptr);
                if (status < 0)
                {
                    LogAndThrowRE("Failed to read register " + register);
                }
                value = Marshal.ReadInt64(value_ptr);
            }
            return (UInt64)value;
        }

        public void SetRegisterUnsafe(int register, RegisterValue value)
        {
            // This is obsolete API, left here only for compatibility
            this.Log(LogLevel.Warning, "Using `SetRegisterUnsafe` API is " +
                     "obsolete. Please change to `SetRegister`.");
            SetRegister(register, value);
        }

        public RegisterValue GetRegisterUnsafe(int register)
        {
            // This is obsolete API, left here only for compatibility
            this.Log(LogLevel.Warning, "Using `GetRegisterUnsafe` API is " +
                     "obsolete. Please change to `GetRegister`.");
            return GetRegister(register);
        }

        public IEnumerable<CPURegister> GetRegisters()
        {
            if (registerMap.Count == 0)
            {
                lock(nativeLock)
                {
                    Int32 num_regs = get_reg_info_size(kelvin_id);
                    for (Int32 i = 0; i < num_regs; i++)
                    {
                        Int32 result = get_reg_info(kelvin_id, i, string_ptr, reg_info_ptr);
                        if (result < 0)
                        {
                            this.ErrorLog("Failed to get register info for index " + i);
                            continue;
                        }
                        var reg_info = Marshal.PtrToStructure<RegInfo>(reg_info_ptr);
                        var cpu_reg = new CPURegister(reg_info.index, reg_info.width,
                                                    reg_info.isGeneral,
                                                    reg_info.isReadOnly);
                        var reg_name = Marshal.PtrToStringAuto(string_ptr);
                        registerMap.Add(reg_info.index, cpu_reg);
                        registerNamesMap.Add(reg_name, reg_info.index);
                    }
                }
            }
            return registerMap.Values.OrderBy(x => x.Index);
        }

        public string[,] GetRegistersValues()
        {
            if (registerMap.Count == 0)
            {
                GetRegisters();
            }
            var result = new Dictionary<string, ulong>();
            foreach (var reg in registerNamesMap)
            {
                Int64 value = 0;
                lock(nativeLock)
                {
                    var status = read_register(kelvin_id, (UInt32)reg.Value, value_ptr);
                    if (status < 0) continue;
                    value = Marshal.ReadInt64(value_ptr);
                }
                result.Add(reg.Key, Convert.ToUInt64(value));
            }
            var table = new Table().AddRow("Name", "Value");
            table.AddRows(result, x => x.Key, x => "0x{0:X}".FormatWith(x.Value));
            return table.ToArray();
        }

        public string CpuLibraryPath
        {
            get
            {
                return cpuLibraryPath;
            }
            set
            {
                if (!String.IsNullOrWhiteSpace(value))
                {
                    if (cpuLibraryRegistered)
                    {
                        this.WarningLog("cpu library already registerd and " +
                                        "should not be updated again.");
                        return;
                    }
                    cpuLibraryPath = value;
                    try
                    {
                        binder = new NativeBinder(this, cpuLibraryPath);
                    }
                    catch (System.Exception e)
                    {
                        LogAndThrowRE("Failed to load CPU library: " + e.Message);
                    }
                    cpuLibraryRegistered = true;

                    this.NoisyLog("Memory block size: {0}, total size: {1}",
                                  dataMemory.SegmentSize, dataMemory.Size);
                    IntPtr[] memoryBlocks = new IntPtr[dataMemory.SegmentCount];
                    for (int i = 0; i < dataMemory.SegmentCount; i++)
                    {
                        this.NoisyLog("Segment {0}: pointer {1:8X}, size {2}",
                                      i, dataMemory.GetSegment(i),
                                      dataMemory.SegmentSize);
                        memoryBlocks[i] = dataMemory.GetSegment(i);
                    }

                    // Simulator initialization code goes here.
                    IntPtr memoryBlocksPtr = Marshal.AllocHGlobal(
                        dataMemory.SegmentCount * Marshal.SizeOf<IntPtr>());
                    Marshal.Copy(memoryBlocks, 0, memoryBlocksPtr, memoryBlocks.Length);
                    lock(nativeLock)
                    {
                        // Initiate a kelvin_sim simulator. For memory access
                        // inside the kelvin_sim, it refers to the memory block
                        // pointers provided by the dataMemory here.
                        //
                        // Note we don't preempt the memory access from other
                        // cores (e.g., MlCoordinator) with the nativeLock. The
                        // control flow we have in CantripOS should not cause
                        // the conflict memory access.
                        kelvin_id = construct_with_memory(maxStringLen,
                                                        (ulong)dataMemory.SegmentSize,
                                                        (ulong)dataMemory.Size,
                                                        memoryBlocksPtr);
                    }
                    Marshal.FreeHGlobal(memoryBlocksPtr);

                    if (kelvin_id < 0)
                    {
                        LogAndThrowRE("Failed to create simulator instance");
                        return;
                    }
                    this.NoisyLog("Construct kelvin");
                }
                else
                {
                    return;
                }
            }
        }

        private void LogAndThrowRE(string info)
        {
            this.Log(LogLevel.Error, info);
            throw new RecoverableException(info);
        }

        protected string cpuLibraryPath;
        private bool cpuLibraryRegistered;

        private readonly MappedMemory dataMemory;
        private NativeBinder binder;
        // lock to ensure only one binding function call is running.
        private readonly object nativeLock = new object();
        private readonly Int32 maxStringLen = 32;
        private IntPtr value_ptr { get; set; }
        private IntPtr error_ptr { get; set; }
        private IntPtr reg_info_ptr { get; set; }
        private IntPtr string_ptr { get; set; }

        [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
        public delegate Int32 ConstructWithMemory(Int32 param0, UInt64 param1, UInt64 param2, IntPtr param3);

        [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
        public delegate Int32 FuncInt32Int32(Int32 param0);

        [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
        public delegate Int32 GetRegInfo(Int32 param0, Int32 param1, IntPtr param2, IntPtr param3);

        [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
        public delegate UInt64 Step(Int32 param0, UInt64 param1, IntPtr param2);

        [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
        public delegate Int32 ReadRegister(Int32 param0, UInt32 param1, IntPtr param2);

        [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
        public delegate Int32 WriteRegister(Int32 param0, Int32 param1, UInt64 param2);

        // Functions from cpuLibrary
        [Import(UseExceptionWrapper = false)]
        // Int32(Int32 max_string_lenth, UInt64 mem_block_size,
        //       UInt64 mem_size, uint8_t ** mem_block_ptr_list);
        private ConstructWithMemory construct_with_memory;
        [Import(UseExceptionWrapper = false)]
        // Int32(Int32 id);
        private FuncInt32Int32 destruct;
        [Import(UseExceptionWrapper = false)]
        // Int32(Int32 id);
        private FuncInt32Int32 reset;
        [Import(UseExceptionWrapper = false)]
        // Int32(Int32 id)
        private FuncInt32Int32 get_reg_info_size;
        [Import(UseExceptionWrapper = false)]
        // Int32(Int32 id, Int32 index, char *name, IntPtr *info);
        private GetRegInfo get_reg_info;
        [Import(UseExceptionWrapper = false)]
        // UInt64(Int32 id, UInt64 step, IntPtr *status);
        private Step step;
        [Import(UseExceptionWrapper = false)]
        // Int32(Int32 id, UInt32 reg_id, IntPtr *value);
        private ReadRegister read_register;
        [Import(UseExceptionWrapper = false)]
        // Int32(Int32 id, Int32 reg_id, UInt64 value);
        private WriteRegister write_register;

        // Register access.
        private Dictionary<Int32, CPURegister> registerMap = new Dictionary<Int32, CPURegister>();
        private Dictionary<string, Int32> registerNamesMap = new Dictionary<string, Int32>();

        // PC magic ID. Defined in
        // https://github.com/riscv/riscv-opcodes/blob/d752f193cec46cfd33300abba2a3a0d020c36755/constants.py#L317
        // It is used to set/get PC by passing this register ID via `read_register`
        // and `write_register`.
        private const int PC_ID = 0x07b1;

        private ulong instructionsExecutedThisRound { get; set; }
        private ulong totalExecutedInstructions { get; set; }

        private Int32 kelvin_id { get; set; }

        // MlTopControlBlock
        private MlTopControlBlock controlBlock;
        private bool controlBlockRegistered = false;
    }


    public class MlTopControlBlock :
            IDoubleWordPeripheral,
            IProvidesRegisterCollection<DoubleWordRegisterCollection>,
            IKnownSize
    {

        public MlTopControlBlock(Machine machine,
                                 KelvinCPU core)
        {
            Machine = machine;
            Core = core;

            HostReqIRQ = new GPIO();
            FinishIRQ = new GPIO();
            InstructionFaultIRQ = new GPIO();

            Core.RegisterControlBlock(this);

            RegistersCollection = new DoubleWordRegisterCollection(this);
            DefineRegisters();

            Reset();
        }

        public void Reset()
        {
            mode = Mode.Freeze | Mode.SwReset;
            RegistersCollection.Reset();
        }

        private void DefineRegisters()
        {
            Registers.IntrState.Define32(this)
                .WithValueField(0, 4,
                                writeCallback: (_, value) =>
                                {
                                    this.NoisyLog("Got {0} to clear IRQ pending bits", value);
                                    irqsPending = irqsPending & ~(InterruptBits)value;
                                    IrqUpdate();
                                },
                                valueProviderCallback: (_) =>
                                {
                                    return (uint)irqsPending;
                                });

            Registers.IntrEnable.Define32(this)
                .WithValueField(0, 4,
                                writeCallback: (_, value) =>
                                {
                                    this.NoisyLog("Got {0} to write IRQ enable bits", value);
                                    irqsEnabled = (InterruptBits)value & InterruptBits.Mask;
                                    IrqUpdate();
                                },
                                valueProviderCallback: (_) =>
                                {
                                    return (uint)irqsEnabled;
                                });

            Registers.IntrTest.Define32(this)
                .WithValueField(0, 4,
                                writeCallback: (_, value) =>
                                {
                                    this.NoisyLog("Got {0} to set IRQ pending bits", value);
                                    irqsPending |= ((InterruptBits)value & InterruptBits.Mask);
                                    IrqUpdate();
                                });

            Registers.Control.Define32(this, resetValue: 0x00000002)
                .WithValueField(0, 24, name: "FREEZE_RESET_PC",
                                writeCallback: (_, val) =>
                                {
                                    Mode newMode = (Mode)val & Mode.Mask;
                                    // Pause the core when either freeze or swreset is asserted.
                                    if ((mode == Mode.Run) && (newMode != Mode.Run))
                                    {
                                        this.NoisyLog("Pausing core.");
                                        Core.IsHalted = true;
                                    }

                                    // Trigger the core's reset when SwReset is deasserted.
                                    if (((mode & Mode.SwReset) != 0) &&
                                        ((newMode & Mode.SwReset) == 0))
                                    {
                                        this.NoisyLog("Resetting core.");
                                        // Reset PC
                                        ulong startAddress = (val >> (int)Mode.NumBits);
                                        if (Core.PC != startAddress)
                                        {
                                            this.DebugLog("Setting PC to 0x{0:X8}.",
                                                          startAddress);
                                            Core.PC = startAddress;
                                        }
                                        Core.Reset();
                                    }

                                    // Raise reset bit, but the core will not be
                                    // reset until SwReset is deasserted.
                                    if ((mode & Mode.SwReset) == 0 &&
                                        ((newMode & Mode.SwReset) != 0))
                                    {
                                        this.NoisyLog("Raise reset bit");
                                    }

                                    // Unpause the core when both freeze and
                                    // SwReset are deasserted.
                                    if ((mode != Mode.Run) && (newMode == Mode.Run))
                                    {
                                        this.NoisyLog("Resuming core.");
                                        Core.IsHalted = false;

                                        Core.Resume();
                                    }

                                    this.mode = newMode;
                                })
                .WithTaggedFlag("VOLT_SEL", 24)
                .WithIgnoredBits(25, 32 - 25);

            // TODO: No implementation to memory bank control.
            Registers.MemoryBankControl.Define32(this)
                .WithTag("MEM_ENABLE", 0, 16)
                .WithIgnoredBits(16, 32 - 16);

            // TODO: No implementation of memory access error.
            Registers.ErrorStatus.Define32(this)
                .WithTaggedFlag("MEM_OUT_OF_RANGE", 0)
                .WithTag("MEM_DISABLE_ACCESS", 1, 8)
                .WithIgnoredBits(9, 32 - 9);

            Registers.InitStart.Define32(this)
                .WithTag("ADDRESS", 0, 22)
                .WithIgnoredBits(22, 32 - 22);

            Registers.InitEnd.Define32(this)
                .WithTag("ADDRESS", 0, 22)
                .WithTaggedFlag("VALID", 22)
                .WithIgnoredBits(23, 32 - 23);

            Registers.InitStatus.Define32(this)
                .WithTaggedFlag("INIT_PENDING", 0)
                .WithTaggedFlag("INIT_DONE", 1)
                .WithIgnoredBits(2, 32 - 2);

        }

        public virtual uint ReadDoubleWord(long offset)
        {
            return RegistersCollection.Read(offset);
        }

        public virtual void WriteDoubleWord(long offset, uint value)
        {
            RegistersCollection.Write(offset, value);
        }

        public void ExecMPause()
        {
            // Pause and trigger a host interrupt indicating completion
            if (mode == Mode.Run)
            {
                this.NoisyLog("Pausing the core for host completion notification.");
            }
            else
            {
                this.ErrorLog("Pausing the core for host completion notification, " +
                              "but core was not expected to be running. Did you " +
                              "clear IsHalted manually?");
            }
            Core.IsHalted = true;
            mode = Mode.Freeze;
            irqsPending |= InterruptBits.Finish;
            IrqUpdate();
        }

        public void ExecEBreak()
        {
            // Pause and trigger a host interrupt indicating completion with
            // fault
            if (mode == Mode.Run)
            {
                this.NoisyLog("Core executed ebreak.");
            }
            else
            {
                this.ErrorLog("Core executed ebreak, but core was not expected " +
                              "to be running. Did you clear IsHalted manually?");
            }
            Core.IsHalted = true;
            mode = Mode.Freeze;
            irqsPending |= InterruptBits.Finish | InterruptBits.InstructionFault;
            IrqUpdate();
        }

        public DoubleWordRegisterCollection RegistersCollection { get; private set; }

        public GPIO HostReqIRQ { get; }
        public GPIO FinishIRQ { get; }
        public GPIO InstructionFaultIRQ { get; }

        private InterruptBits irqsEnabled;
        private InterruptBits irqsPending;

        private void IrqUpdate()
        {
            InterruptBits irqsPassed = irqsEnabled & irqsPending;
            HostReqIRQ.Set((irqsPassed & InterruptBits.HostReq) != 0);
            FinishIRQ.Set((irqsPassed & InterruptBits.Finish) != 0);
            InstructionFaultIRQ.Set((irqsPassed & InterruptBits.InstructionFault) != 0);
        }

        private Mode mode;
        private readonly Machine Machine;
        private readonly KelvinCPU Core;
        // Length of register space. Required by IKnowSize.
        public long Size => 0x1000;
        private enum Registers
        {
            IntrState = 0x00,
            IntrEnable = 0x04,
            IntrTest = 0x08,
            Control = 0x0C,
            MemoryBankControl = 0x10,
            ErrorStatus = 0x14,
            InitStart = 0x18,
            InitEnd = 0x1C,
            InitStatus = 0x20,
        };
        [Flags]
        private enum Mode
        {
            Run = 0x00,
            Freeze = 0x01,
            SwReset = 0x02,
            Mask = 0x03,
            NumBits = 2,
        };
        [Flags]
        private enum InterruptBits
        {
            HostReq = 1,
            Finish = 2,
            InstructionFault = 4,
            Mask = 7,
        };
    }

}