#region License
/* 
 * Copyright (C) 1999-2024 John Källén.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2, or (at your option)
 * any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; see the file COPYING.  If not, write to
 * the Free Software Foundation, 675 Mass Ave, Cambridge, MA 02139, USA.
 */
#endregion

using Reko.Core;
using Reko.Core.Code;
using Reko.Core.Diagnostics;
using Reko.Core.Expressions;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;

namespace Reko.Scanning
{
    public class PromoteBlockWorkItem : WorkItem
    {
        private static readonly TraceSwitch trace = new TraceSwitch(nameof(PromoteBlockWorkItem), "Trace the workings of PromoteBlockWorkItem") { Level = TraceLevel.Info };

        public Block Block;
        public Procedure ProcNew;
        public IScannerServices  Scanner;
        public Program Program;

        public PromoteBlockWorkItem(
            Address addr,
            IScannerServices scanner,
            Program program,
            Block  block,
            Procedure procNew) : base(addr)
        {
            Scanner = scanner;
            Program = program;
            Block = block;
            ProcNew = procNew;
        }

        public override void Process()
        {
            var movedBlocks = new HashSet<Block>();
            var stack = new Stack<IEnumerator<Block>>();
            stack.Push(new Block[] { Block }.Cast<Block>().GetEnumerator());
            var replacers = new Dictionary<Procedure, IdentifierRelocator>();
            while (stack.Count != 0)
            {
                DumpBlocks(Block.Procedure);
                var e = stack.Peek();
                if (!e.MoveNext())
                {
                    stack.Pop();
                    continue;
                }
                var b = e.Current;
                if (b.Procedure == ProcNew || b == b.Procedure.ExitBlock || b.Procedure.EntryBlock.Succ[0] == b)
                    continue;

                trace.Verbose("PBW:     Visiting block {0}, stack depth {1}", b.DisplayName, stack.Count);
                if (!replacers.TryGetValue(b.Procedure, out var replacer))
                {
                    replacer = new IdentifierRelocator(b.Procedure.Frame, ProcNew.Frame);
                    replacers.Add(b.Procedure, replacer);
                }
                b.Procedure.RemoveBlock(b);
                ProcNew.AddBlock(b);
                b.Procedure = ProcNew;
                movedBlocks.Add(b);
                foreach (var stm in b.Statements)
                {
                    stm.Instruction = replacer.ReplaceIdentifiers(stm.Instruction);
                }
                if (b.Succ.Count > 0)
                    stack.Push(b.Succ.GetEnumerator());
            }
            foreach (var b in movedBlocks)
            {
                FixExitEdges(b);
                FixInboundEdges(b);
                FixOutboundEdges(b);
               //  SanityCheck(b);
            }
        }

        [Conditional("DEBUG")]
        public static void SanityCheck(Block block)
        {
         //   Debug.Assert(block.Pred.Count == 0 && block != block.Procedure.EntryBlock);
        }

        [Conditional("DEBUG")]
        private void DumpBlocks(Procedure procedure)
        {
            trace.Verbose("{0}", procedure.Name);
            foreach (var block in procedure.ControlGraph.Blocks)
            {
                trace.Verbose("  {0}; {1}", block.DisplayName, block.Procedure.Name);
            }
        }

        public void FixInboundEdges(Block blockToPromote)
        {
            trace.Verbose("PBW: Fixing inbound edges of {0}", blockToPromote.DisplayName);

            // Get all blocks that are from "outside" blocks.
            var inboundBlocks = blockToPromote.Pred.Where(p => p.Procedure != ProcNew).ToArray();
            foreach (var inb in inboundBlocks)
            {
                if (inb.Statements.Count > 0)
                {
                    var lastAddress = GetAddressOfLastInstruction(inb);
                    var state = inb.Procedure.Architecture.CreateProcessorState();
                    var callRetThunkBlock = Scanner.CreateCallRetThunk(lastAddress, inb.Procedure, state, ProcNew);
                    ReplaceSuccessorsWith(inb, blockToPromote, callRetThunkBlock);
                    callRetThunkBlock.Pred.Add(inb);
                }
                else
                {
                    var stmLast = inb.Statements.Add(
                        inb.Address,
                        new CallInstruction(
                            new ProcedureConstant(Program.Platform.PointerType, ProcNew),
                            new CallSite(0, 0)));
                    Program.CallGraph.AddEdge(stmLast, ProcNew);
                    inb.Statements.Add(inb.Address, new ReturnInstruction());
                    inb.Procedure.ControlGraph.AddEdge(inb, inb.Procedure.ExitBlock);
                }
            }
            foreach (var p in inboundBlocks)
            {
                blockToPromote.Pred.Remove(p);
            }
        }

        private Address GetAddressOfLastInstruction(Block inboundBlock)
        {
            if (inboundBlock.Statements.Count == 0)
                return Program.Platform.MakeAddressFromLinear(0, true);
            var stmLast = inboundBlock.Statements[^1];
            return inboundBlock.Address != null
                ? inboundBlock.Address + (stmLast.Address - inboundBlock.Statements[0].Address)
                : stmLast.Address;
        }

        public void FixOutboundEdges(Block block)
        {
            trace.Verbose("PBW: Fixing outbound edges of {0}", block.DisplayName);
            for (int i = 0; i < block.Succ.Count; ++i)
            {
                var s = block.Succ[i];
                if (s.Procedure == block.Procedure)
                    continue;
                if (s.Procedure.EntryBlock.Succ[0] == s)
                {
                    // s is the first block of a (different) procedure
                    var lastAddress = GetAddressOfLastInstruction(block);
                    var state = block.Procedure.Architecture.CreateProcessorState();
                    var retCallThunkBlock = Scanner.CreateCallRetThunk(lastAddress, block.Procedure, state, s.Procedure);
                    block.Succ[i] = retCallThunkBlock;
                    retCallThunkBlock.Pred.Add(block);
                    s.Pred.Remove(block);
                }
            }
        }

        private static void ReplaceSuccessorsWith(Block block, Block blockOld, Block blockNew)
        {
            for (int s = 0; s < block.Succ.Count; ++s)
            {
                if (block.Succ[s] == blockOld)
                    block.Succ[s] = blockNew;
            }
        }

        private void FixExitEdges(Block block)
        {
            trace.Verbose("PBW: Fixing exit edges of {0}", block.DisplayName);
            for (int i = 0; i < block.Succ.Count; ++i)
            {
                var s = block.Succ[i];
                if (s.Procedure != ProcNew && s == s.Procedure.ExitBlock)
                {
                    s.Pred.Remove(block);
                    ProcNew.ExitBlock.Pred.Add(block);
                    block.Succ[i] = ProcNew.ExitBlock;
                }
            }
        }
    }
}
